Binary Heap Operations

In the context of a binary heap, “sift” means to move an element up or down the heap until it satisfies the heap property, which is that the value of each node must be greater than or equal to its children in the case of a max heap, or less than or equal to its children in the case of a min heap. The _sift_up method moves an element up the heap until it is greater than or equal to its parent, and the _sift_down method moves an element down the heap until it is less than or equal to its children.

• The below heap has a compare method which decides number is big or small, this functionality is to decide if it will be a Min heap or Max heap.

class BinaryHeap:
    """
    A binary heap data structure.
    """
    def __init__(self, is_max_heap=True):
        """
        Initializes an empty binary heap.
 
        Args:
        - is_max_heap: a boolean indicating whether the heap should be a max heap (True)
                       or a min heap (False).
        """
        self.heap = []
        self.is_max_heap = is_max_heap
 
    def __len__(self):
        """
        Returns the number of elements in the heap.
        """
        return len(self.heap)
 
    def __bool__(self):
        """
        Returns True if the heap is not empty, False otherwise.
        """
        return bool(self.heap)
 
    def _compare(self, a, b):
        """
        Compares two elements a and b based on whether the heap is a max heap or a min heap.
 
        Returns True if a is greater than or equal to b in a max heap,
        or a is less than or equal to b in a min heap.
        """
        return a >= b if self.is_max_heap else a <= b
 
    def _swap(self, i, j):
        """
        Swaps the elements at indices i and j in the heap.
        """
        self.heap[i], self.heap[j] = self.heap[j], self.heap[i]
 
    def _sift_up(self, index):
        """
        Sifts the element at the given index up in the heap until it satisfies the heap property.
        """
        while index > 0:
            parent = (index - 1) // 2
            if self._compare(self.heap[index], self.heap[parent]):
                self._swap(index, parent)
                index = parent
            else:
                break
 
    def _sift_down(self, index):
        """
        Sifts the element at the given index down in the heap until it satisfies the heap property.
        """
        while index < len(self):
            left_child = 2 * index + 1
            right_child = 2 * index + 2
 
            if left_child >= len(self):
                break
 
            min_max_child = left_child
            if right_child < len(self) and self._compare(self.heap[right_child], self.heap[left_child]):
                min_max_child = right_child
 
            if self._compare(self.heap[min_max_child], self.heap[index]):
                self._swap(index, min_max_child)
                index = min_max_child
            else:
                break
 
    def insert(self, element):
        """
        Inserts an element into the heap.
 
        Args:
        - element: the element to insert.
        """
        self.heap.append(element)
        self._sift_up(len(self) - 1)
 
    def extract(self):
        """
        Extracts the minimum or maximum element from the heap, depending on whether
        the heap is a min heap or a max heap.
 
        Returns:
        - the minimum or maximum element from the heap, or None if the heap is empty.
        """
        if not self:
            return None
 
        min_max_element = self.heap[0]
        last_element = self.heap.pop()
 
        if self:
            self.heap[0] = last_element
            self._sift_down(0)
 
        return min_max_element
 
# Create a max heap
max_heap = BinaryHeap(is_max_heap=True)
 
# Insert some elements
max_heap.insert(5)
max_heap.insert(3)
max_heap.insert(8)
max_heap.insert(4)
 
# Extract the maximum element
max_element = max_heap.extract()
print(max_element)  # Output: 8
 
# Create a min heap
min_heap = BinaryHeap(is_max_heap=False)
 
# Insert some elements
min_heap.insert(5)
min_heap.insert(3)
min_heap.insert(8)
min_heap.insert(4)
 
# Extract the minimum element
min_element = min_heap.extract()
print(min_element)  # Output: 3