Inspired by exercise on HackerRank where I'm expected to implement binary heap with additional method that allows to remove any value from heap I decided to clean up and expand my code a little. Expansion is an ability to pass custom comparator to the heap, and method to remove first element from heap (which for some reason was missing from the exercise).
Heap takes integer values, like in original exercise, but all it needs to be done to change that is to change typing and pass right comparators. Original problem statement assumed that unique integers, but the heap should work also for repeated values.
I've passed my code through PyCharm and removed issues it detected, but I want to know if I'm following other best python practices. If my algorithm is flawed I would also like to know that. I know PEP8 suggests maximum line length of 80, but I'm more comfortable with PyCharm default limit of 120. (In my current workplace limit is even higher.)
import unittest
import doctest
from typing import List, Optional, Callable
class Heap:
"""
Binary heap. Array implementation.
Additional function for erasing element of certain value
>>> heap = Heap()
>>> heap.add(4)
>>> heap.add(2)
>>> heap.add(3)
>>> heap.add(1)
>>> print(heap.erase_value(17))
False
>>> print(heap.erase_value(2))
True
>>> print(heap.pop_first())
1
>>> print(heap.pop_first())
3
>>> print(heap.peek_first())
4
>>> print(heap.pop_first())
4
>>> print(heap.pop_first())
None
"""
array: List[int]
comparator: Callable[[int, int], int]
def __init__(self, comparator: Callable[[int, int], int] = lambda a, b: -1 if a < b else 0 if a == b else 1):
"""
Array initiated with guard, so indexing starts with 1, which makes getting children/parent easier.
:param comparator: Comparator to use. Default results in min-heap.
"""
self.array = [None]
self.comparator = comparator
def _left_child_index(self, parent_index: int) -> Optional[int]:
"""
Give the index of the left child of the parent.
:param parent_index: Index of the parent.
:return: Index of the left child or None if parent does not have left child.
"""
child_index = parent_index * 2
if child_index < len(self.array):
return child_index
return None
def _right_child_index(self, parent_index: int) -> Optional[int]:
"""
Give the index of the right child of the parent.
:param parent_index: Index of the parent.
:return: Index of the right child or None if parent does not have right child.
"""
child_index = parent_index * 2 + 1
if child_index < len(self.array):
return child_index
return None
@staticmethod
def _parent_index(child_index: int) -> Optional[int]:
"""
Get the index of the parent of the child.
:param child_index: Index of the child.
:return: Index of the parent or None if child is root.
"""
if child_index <= 1:
return None
return child_index // 2
def _swap_indexes(self, index1: int, index2: int) -> None:
"""
Swap values from two indexes.
:param index1: Index of first element.
:param index2: Index of second element.
"""
tmp = self.array[index1]
self.array[index1] = self.array[index2]
self.array[index2] = tmp
def _heapify_up(self, index: int) -> None:
"""
Update heap tree by floating current element up till it's in the right place.
:param index: Current element index.
"""
parent_index = self._parent_index(index)
if parent_index is None:
return
if self.comparator(self.array[index], self.array[parent_index]) < 0:
self._swap_indexes(index, parent_index)
self._heapify_up(parent_index)
def _heapify_down(self, index: int) -> None:
"""
Update heap tree by dropping current element down till it's in the right place.
:param index: Current element index.
"""
if index >= len(self.array):
return
left_child = self._left_child_index(index)
right_child = self._right_child_index(index)
if left_child is None:
return
if self.comparator(self.array[left_child], self.array[index]) < 0:
if (right_child is not None) and self.comparator(self.array[right_child], self.array[left_child]) < 0:
self._swap_indexes(index, right_child)
self._heapify_down(right_child)
else:
self._swap_indexes(index, left_child)
self._heapify_down(left_child)
else:
if (right_child is not None) and self.comparator(self.array[right_child], self.array[index]) < 0:
self._swap_indexes(index, right_child)
self._heapify_down(right_child)
def add(self, val: int) -> None:
"""
Add value to the heap.
:param val: Value to add.
"""
self.array.append(val)
self._heapify_up(len(self.array) - 1)
def erase_value(self, val: int) -> bool:
"""
Erase first instance of value from the heap.
:param val: Value to erase.
:return: True if something was erased, false otherwise.
"""
try:
index = self.array.index(val)
except ValueError:
return False
self._swap_indexes(index, len(self.array) - 1)
self.array.pop()
parent_index = self._parent_index(index)
if index >= len(self.array):
return True
if (parent_index is None) or self.comparator(self.array[parent_index], self.array[index]):
self._heapify_down(index)
else:
self._heapify_up(index)
return True
def pop_first(self) -> Optional[int]:
"""
Erase and get the first element from the heap.
:return: The first element from the heap or None if heap is empty.
"""
first_element = self.peek_first()
if first_element is None:
return None
self.erase_value(first_element)
return first_element
def peek_first(self) -> Optional[int]:
"""
Get the first element from the heap.
:return: The first element from the heap or None if heap is empty.
"""
if len(self.array) <= 1:
return None
return self.array[1]
class TestHeap(unittest.TestCase):
heap_min: Heap
heap_max: Heap
def setUp(self):
self.heap_min = Heap()
self.heap_max = Heap(lambda a, b: -1 if a > b else 0 if a == b else 1)
def test_inserting_in_sorted_order(self):
for num in range(1, 16):
self.heap_min.add(num)
for num in range(1, 16):
self.assertEqual(self.heap_min.pop_first(), num)
assert self.heap_min.peek_first() is None
def test_inserting_in_reverse_order(self):
for num in range(15, 0, -1):
self.heap_min.add(num)
for num in range(1, 16):
self.assertEqual(self.heap_min.pop_first(), num)
assert self.heap_min.peek_first() is None
def test_removing_middle_first(self):
for num in range(1, 16):
self.heap_min.add(num)
for num in range(6, 11):
self.assertTrue(self.heap_min.erase_value(num))
for num in range(1, 6):
self.assertEqual(self.heap_min.pop_first(), num)
for num in range(11, 16):
self.assertEqual(self.heap_min.pop_first(), num)
assert self.heap_min.peek_first() is None
def test_erasing_nonexistent(self):
for num in range(1, 16):
self.heap_min.add(num)
self.assertFalse(self.heap_min.erase_value(17))
self.assertFalse(self.heap_min.erase_value(42))
self.assertFalse(self.heap_min.erase_value(-1))
self.assertTrue(self.heap_min.erase_value(2))
self.assertFalse(self.heap_min.erase_value(2))
def test_non_unique_values(self):
for num in range(1, 8):
self.heap_min.add(num)
for num in range(1, 8):
self.heap_min.add(num)
self.assertTrue(self.heap_min.erase_value(1))
self.assertTrue(self.heap_min.erase_value(1))
self.assertFalse(self.heap_min.erase_value(1))
self.assertTrue(self.heap_min.erase_value(4))
self.assertTrue(self.heap_min.erase_value(4))
for num in range(2, 8):
if num == 4:
continue
self.assertEqual(self.heap_min.pop_first(), num)
self.assertEqual(self.heap_min.pop_first(), num)
assert self.heap_min.peek_first() is None
def test_max_comparator(self):
for num in range(1, 16):
self.heap_max.add(num)
for num in range(15, 0, -1):
self.assertEqual(self.heap_max.pop_first(), num)
assert self.heap_max.peek_first() is None
if __name__ == '__main__':
doctest.testmod()
unittest.main()
