Python 3 implementation of binary heap

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()
>>> 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:
"""

"""
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):
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):
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):
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.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):
for num in range(1, 8):
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):
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()


Excellent

• Type annotations help the IDE make sense of dynamically typed code.
• Doctest examples show people how to use your code with theirs.
• Unit tests ensure that future changes are less likely to break functionality.

Naming issues

• When doing the opposite of an action I would use the antonym of the method name. So in this case I would rename erase_value to delete, remove or another antonym of "add".
• The _first suffix looks redundant - pop and peek are common names for what you are doing.
• A common pattern for first and second parameters like in _swap_indexes or your lambdas is simply first and second. These are lexicographically more distinct than index1 and index2, so it's much harder to confuse the two.

Keeping it simple

• By all means write comments to clarify things to yourself while programming, but then see if you can simplify the code so much that you can remove it. Comments should explain the "why", not the "how", "what" or "when."
• I would pull out both lambdas and give them names. That way the "Default results in min-heap." comment becomes redundant and -1 if a < b else 0 if a == b else 1 can be much more legible as multiple lines.
• Unless binary heaps are special in some way starting indexing at 1 will be extremely confusing to most programmers. Python as a rule does not do 1-based indexing.

Tests

• Testing should be done by an external process, not every time the program runs.
• Test naming is tricky. I know almost every guide out there uses the standard "test [function name]" or "test [high level requirement]" naming scheme, but there's a couple tricks to make them much more useful:
• Start the test name with "should." (In Python you need to configure your test framework to look for should_ or just use test_should_.) This forces you to think about the goal of the test, to prove that the code does what it should do rather than just what sort of data or function is involved in the test. For example, renaming test_erasing_nonexistent to should_return_false_when_removing_nonexistent_number would clarify what the test is trying to prove.
• I can tell from the tests that they were written after the code, because I used to write tests the same way. With TDD you would discover that there is no point testing 16 values when one, two or three would do to prove the point. This would also make the tests much easier to read, because as a reviewer I would be able to verify with the absolute minimal cognitive load that they do what they say.
• Each test should ideally contain only one assertion. (The only exception to this that I've found is when a guard assertion results in more useful error messages, such as self.assertEqual(200, request.status_code, request.data).) While resulting in lots more test methods this gives several advantages:
• Each test is much easier to read.
• It's much easier to spot redundant, irrelevant or wrong tests.
• You won't have to play whack-a-mole with assertions to get to the end of the test.

Mypy by default does very few checks; last time I configured it I ended up with this:

[mypy]
check_untyped_defs = true
disallow_untyped_defs = true
ignore_missing_imports = true
no_implicit_optional = true
warn_redundant_casts = true
warn_return_any = true
warn_unused_ignores = true


This would flag some more issues, like __init__ and test methods missing return types (which should be None).

• Thanks, it helped a lot. One thing I'm not sure is making methods that are implementation dependent public. Wouldn't it be better to keep public just methods that are meaningful in all implementations of binary heap (or maybe other heaps too) in case I'll change implementation in the future? Nov 4, 2018 at 17:43