minheap

priority_queue.py

class PQ:
    """
    Priority queue implemented as a heap (minheap).
    """

    def __init__(self ):
        self.heap = [0]
        self.size = 0

    def insert(self, x):
        self.heap.append(x)
        self.size += 1
        self._bubble_up(self.size)

    def pop_min(self):
        self._swap(1, self.size)
        ret = self.heap.pop()
        self.size -= 1
        self._bubble_down(1)
        return ret

    def build_heap(self, A):
        """ Heapsort array A """
        self.size = len(A)
        i = self.size // 2
        self.heap = [0] + list(A)
        while i > 0:
            self._bubble_down(i)
            i -= 1

    def _bubble_up(self, i):
        parent = i // 2
        while parent:
            if self.heap[parent] > self.heap[i]:
                self._swap(i, parent)
                i = parent
                parent = i // 2
            else:
                return

    def _bubble_down(self, i):
        while i*2 <= self.size:
            mc = self._min_child(i)
            if self.heap[mc] < self.heap[i]:
                self._swap(mc, i)
                i = mc
            else:
                return

    def _min_child(self, i):
        if i * 2 <= self.size:
            if self.size < i * 2 + 1:
                return i*2
            l, r = self.heap[i*2], self.heap[i*2+1]
            if l < r:
                return i * 2
            else:
                return i * 2 + 1
        return i

    def _swap(self, n1, n2):
        tmp = self.heap[n1]
        self.heap[n1] = self.heap[n2]
        self.heap[n2] = tmp