heap_generic.py

fromcollections.abcimportCallable


classHeap:
"""
 A generic Heap class, can be used as min or max by passing the key function
 accordingly.
 """

def__init__(self, key: Callable|None=None) ->None:
# Stores actual heap items.
self.arr: list= []
# Stores indexes of each item for supporting updates and deletion.
self.pos_map: dict= {}
# Stores current size of heap.
self.size=0
# Stores function used to evaluate the score of an item on which basis ordering
# will be done.
self.key=keyor (lambdax: x)

def_parent(self, i: int) ->int|None:
"""Returns parent index of given index if exists else None"""
returnint((i-1) /2) ifi>0elseNone

def_left(self, i: int) ->int|None:
"""Returns left-child-index of given index if exists else None"""
left=int(2*i+1)
returnleftif0<left<self.sizeelseNone

def_right(self, i: int) ->int|None:
"""Returns right-child-index of given index if exists else None"""
right=int(2*i+2)
returnrightif0<right<self.sizeelseNone

def_swap(self, i: int, j: int) ->None:
"""Performs changes required for swapping two elements in the heap"""
# First update the indexes of the items in index map.
self.pos_map[self.arr[i][0]], self.pos_map[self.arr[j][0]] = (
self.pos_map[self.arr[j][0]],
self.pos_map[self.arr[i][0]],
 )
# Then swap the items in the list.
self.arr[i], self.arr[j] =self.arr[j], self.arr[i]

def_cmp(self, i: int, j: int) ->bool:
"""Compares the two items using default comparison"""
returnself.arr[i][1] <self.arr[j][1]

def_get_valid_parent(self, i: int) ->int:
"""
 Returns index of valid parent as per desired ordering among given index and
 both it's children
 """
left=self._left(i)
right=self._right(i)
valid_parent=i

ifleftisnotNoneandnotself._cmp(left, valid_parent):
valid_parent=left
ifrightisnotNoneandnotself._cmp(right, valid_parent):
valid_parent=right

returnvalid_parent

def_heapify_up(self, index: int) ->None:
"""Fixes the heap in upward direction of given index"""
parent=self._parent(index)
whileparentisnotNoneandnotself._cmp(index, parent):
self._swap(index, parent)
index, parent=parent, self._parent(parent)

def_heapify_down(self, index: int) ->None:
"""Fixes the heap in downward direction of given index"""
valid_parent=self._get_valid_parent(index)
whilevalid_parent!=index:
self._swap(index, valid_parent)
index, valid_parent=valid_parent, self._get_valid_parent(valid_parent)

defupdate_item(self, item: int, item_value: int) ->None:
"""Updates given item value in heap if present"""
ifitemnotinself.pos_map:
return
index=self.pos_map[item]
self.arr[index] = [item, self.key(item_value)]
# Make sure heap is right in both up and down direction.
# Ideally only one of them will make any change.
self._heapify_up(index)
self._heapify_down(index)

defdelete_item(self, item: int) ->None:
"""Deletes given item from heap if present"""
ifitemnotinself.pos_map:
return
index=self.pos_map[item]
delself.pos_map[item]
self.arr[index] =self.arr[self.size-1]
self.pos_map[self.arr[self.size-1][0]] =index
self.size-=1
# Make sure heap is right in both up and down direction. Ideally only one
# of them will make any change- so no performance loss in calling both.
ifself.size>index:
self._heapify_up(index)
self._heapify_down(index)

definsert_item(self, item: int, item_value: int) ->None:
"""Inserts given item with given value in heap"""
arr_len=len(self.arr)
ifarr_len==self.size:
self.arr.append([item, self.key(item_value)])
else:
self.arr[self.size] = [item, self.key(item_value)]
self.pos_map[item] =self.size
self.size+=1
self._heapify_up(self.size-1)

defget_top(self) ->tuple|None:
"""Returns top item tuple (Calculated value, item) from heap if present"""
returnself.arr[0] ifself.sizeelseNone

defextract_top(self) ->tuple|None:
"""
 Return top item tuple (Calculated value, item) from heap and removes it as well
 if present
 """
top_item_tuple=self.get_top()
iftop_item_tuple:
self.delete_item(top_item_tuple[0])
returntop_item_tuple


deftest_heap() ->None:
"""
 >>> h = Heap() # Max-heap
 >>> h.insert_item(5, 34)
 >>> h.insert_item(6, 31)
 >>> h.insert_item(7, 37)
 >>> h.get_top()
 [7, 37]
 >>> h.extract_top()
 [7, 37]
 >>> h.extract_top()
 [5, 34]
 >>> h.extract_top()
 [6, 31]
 >>> h = Heap(key=lambda x: -x) # Min heap
 >>> h.insert_item(5, 34)
 >>> h.insert_item(6, 31)
 >>> h.insert_item(7, 37)
 >>> h.get_top()
 [6, -31]
 >>> h.extract_top()
 [6, -31]
 >>> h.extract_top()
 [5, -34]
 >>> h.extract_top()
 [7, -37]
 >>> h.insert_item(8, 45)
 >>> h.insert_item(9, 40)
 >>> h.insert_item(10, 50)
 >>> h.get_top()
 [9, -40]
 >>> h.update_item(10, 30)
 >>> h.get_top()
 [10, -30]
 >>> h.delete_item(10)
 >>> h.get_top()
 [9, -40]
 """


if__name__=="__main__":
importdoctest

doctest.testmod()