Skip to content

Instantly share code, notes, and snippets.

@roma-glushko
Last active May 26, 2021 08:22
Show Gist options
  • Save roma-glushko/43678fcec904209adb276f33a56ce9de to your computer and use it in GitHub Desktop.
Save roma-glushko/43678fcec904209adb276f33a56ce9de to your computer and use it in GitHub Desktop.
Implementation of heap (priority queue) from my blog post: https://www.romaglushko.com/blog/heapify/
from typing import List
class PriorityQueue:
"""
Represents the heap and preserves the heap property during adding/removing elements
"""
items: List[int]
def __init__(self, items: List[int]):
self.items = self.build_heap(items)
def build_heap(self, items: List[int]) -> List[int]:
"""
Turn an unsorted array into a heap
"""
items_count = len(items)
for i in range(items_count // 2, -1, -1):
items = self.heapify(items, i)
return items
def heapify(self, items: List[int], node_idx: int, root_idx: int = 0) -> List[int]:
"""
Check and fix violations of the heap property recursively
"""
items_count = len(items)
largest_idx = node_idx
# formulas for zero-indexed arrays
left_child_idx = 2 * (node_idx - root_idx) + 1 + root_idx
right_child_idx = 2 * (node_idx - root_idx) + 2 + root_idx
# is the left child node bigger than parent node?
if left_child_idx < items_count and items[left_child_idx] > items[largest_idx]:
largest_idx = left_child_idx
# is the right child node bigger than parent or left node?
if right_child_idx < items_count and items[right_child_idx] > items[largest_idx]:
largest_idx = right_child_idx
# let's fix a violation of the heap property
if largest_idx != node_idx:
items[node_idx], items[largest_idx] = items[largest_idx], items[node_idx]
return self.heapify(items, largest_idx, root_idx=root_idx)
return items
def push(self, item: int):
self.items.append(item)
idx = len(self.items) - 1
parent_idx = idx // 2
if idx % 2 == 0:
# calculating parents in zero-indexed array
parent_idx -= 1
while idx > 0 and self.items[idx] > self.items[parent_idx]:
self.items[parent_idx], self.items[idx] = self.items[idx], self.items[parent_idx]
idx = parent_idx
parent_idx = idx // 2
if idx % 2 == 0:
# calculating parents in zero-indexed array
parent_idx -= 1
def pop(self):
"""
Extract the next item from the heap according to priority (the next biggest element in our max heap case)
"""
item = self.items[0]
items_count = len(self.items)
self.items[0] = self.items[items_count - 1] # replace extracted max element with one of the balanced leaves
del self.items[items_count - 1]
self.items = self.heapify(self.items, 0)
return item
def size(self) -> int:
return len(self.items)
def reverse_sort(self) -> List[int]:
items = self.items
items_count = len(items)
for i in range(items_count - 1, 0, -1):
# move the current max element to the end of the reduced array
items[i], items[0] = items[0], items[i]
# find the next max element/restore the heap property
items = self.heapify(items, 0)
return items
def sort(self) -> List[int]:
"""
In-place version of the heap sort.
It breaks the heap property, so be aware of that.
"""
items = self.items
items_count = len(items)
for i in range(1, items_count - 1):
# find the next max element/restore the heap property
items = self.heapify(items, i, root_idx=i)
return items
def heap_sort(heap: PriorityQueue) -> List[int]:
n = heap.size()
result = []
for i in range(n):
result.append(heap.pop())
return result
if __name__ == "__main__":
heap = PriorityQueue([2, 7, 26, 25, 19, 17, 1, 90, 3, 36])
print(heap_sort(heap))
heap = PriorityQueue([2, 7, 26, 25, 19, 17, 1, 90, 3, 36])
print(heap.reverse_sort())
heap = PriorityQueue([2, 7, 26, 25, 19, 17, 1, 90, 3, 36])
print(heap.sort())
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment