heap.py

class MinHeap:
    def __init__(self, data=None):
        """
        Initialize a Min Heap.
        :param data: Optional list to build the heap from.
        """
        self.heap = data if data else []
        if data:
            self.build_min_heap()
    
    def _parent(self, index):
        """Get parent index using bitwise right shift (equivalent to index // 2)."""
        return (index - 1) >> 1
    
    def _left_child(self, index):
        """Get left child index using bitwise left shift (equivalent to 2 * index + 1)."""
        return (index << 1) + 1
    
    def _right_child(self, index):
        """Get right child index using bitwise left shift (equivalent to 2 * index + 2)."""
        return (index << 1) + 2
    
    def heapify(self, index):
        """
        Heapify the subtree rooted at index.
        Ensures the heap property is maintained.
        """
        smallest = index
        left = self._left_child(index)
        right = self._right_child(index)
        
        if left < len(self.heap) and self.heap[left] < self.heap[smallest]:
            smallest = left
        if right < len(self.heap) and self.heap[right] < self.heap[smallest]:
            smallest = right
        if smallest != index:
            self.heap[index], self.heap[smallest] = self.heap[smallest], self.heap[index]
            self.heapify(smallest)
    
    def build_min_heap(self):
        """
        Convert an unordered list into a min heap.
        """
        for i in range((len(self.heap) - 2) >> 1, -1, -1):
            self.heapify(i)
    
    def push(self, item):
        """
        Insert a new element into the heap.
        """
        self.heap.append(item)
        current = len(self.heap) - 1
        while current > 0 and self.heap[current] < self.heap[self._parent(current)]:
            parent = self._parent(current)
            self.heap[current], self.heap[parent] = self.heap[parent], self.heap[current]
            current = parent
    
    def pop(self):
        """
        Remove and return the smallest element (root) from the heap.
        """
        if not self.heap:
            raise IndexError("Heap is empty")
        if len(self.heap) == 1:
            return self.heap.pop()
        root = self.heap[0]
        self.heap[0] = self.heap.pop()
        self.heapify(0)
        return root
    
    def __str__(self):
        return str(self.heap)

# Testing the MinHeap implementation
# Test build_min_heap
heap = MinHeap([9, 5, 6, 2, 3, 8])
assert heap.heap[0] == 2, "build_min_heap failed"

# Test push
heap.push(1)
assert heap.heap[0] == 1, "push failed"

# Test pop
min_val = heap.pop()
assert min_val == 1, "pop failed"
assert 1 not in heap.heap, "pop did not remove the minimum value"

# Test heapify
heap = MinHeap([10, 1, 3])
heap.heapify(0)
assert heap.heap[0] == 1, "heapify failed"

# Test pop on empty heap
empty_heap = MinHeap()
try:
    empty_heap.pop()
    assert False, "pop on empty heap should raise IndexError"
except IndexError:
    pass  # Expected behavior

print("All tests passed successfully!")