Search code examples
pythonmin-heapmax-heap

MinHeap/MaxHeap implementation in Python


I scoured the web to come up with this implementation of MinHeap and Maxheap in Python.

import heapq


class MinHeap:
    def __init__(self):
        self.heap = []

    def push(self, item):
        heapq.heappush(self.heap, item)

    def pop(self):
        return heapq.heappop(self.h)

    def __getitem__(self, item):
        return self.heap[item]

    def __len__(self):
        return len(self.h)


class MaxHeap(MinHeap):
    def push(self, item):
        heapq.heappush(self.heap, Comparator(item))

    def pop(self):
        return heapq.heappop(self.h)

    def __getitem__(self, i):
        return self.heap[i].val


class Comparator:
    def __init__(self, val):
        self.val = val

    def __lt__(self, other):
        return self.val > self.other

    def __eq__(self, other):
        return self.val == self.other

    def __str__(self):
        return str(self.val)

Now I need to add a peek method to both these classes. In it's current implementation I could pop and push back. But my question is, is there a better way to do it. Something in O(1) time.


Solution

  • These classes are based on Python's heapq structure, which is built on a standard Python list. The smallest item in the minheap and the largest item in the maxheap is at index zero. So just return

    self.heap[0]
    

    If the heap is empty, that will cause an error, but that probably is what you want.