Leetcode 295 Solution

This article provides solution to leetcode question 295 (find-median-from-data-stream)
Categories:
Tags:
Access this page by simply typing in "lcs 295" in your browser address bar if you have bunnylol configured.

https://leetcode.com/problems/find-median-from-data-stream

Solution

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

    def heappush(self, val):
        heapq.heappush(self.heap, val)

    def heappop(self):
        return heapq.heappop(self.heap)

    def top(self):
        return self.heap[0]

    def len(self):
        return len(self.heap)


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

    def heappush(self, val):
        heapq.heappush(self.heap, -val)

    def heappop(self):
        return -heapq.heappop(self.heap)

    def top(self):
        return -self.heap[0]

    def len(self):
        return len(self.heap)


class MedianFinder:

    def __init__(self):
        """
        initialize your data structure here.
        """
        self.heap1 = MaxHeap()
        self.heap2 = MinHeap()

    def addNum(self, num: int) -> None:
        if self.heap1.len() and num < self.heap1.top():
            self.heap1.heappush(num)
        else:
            self.heap2.heappush(num)

        if self.heap1.len() == self.heap2.len() + 2:
            self.heap2.heappush(self.heap1.heappop())
        elif self.heap2.len() == self.heap1.len() + 2:
            self.heap1.heappush(self.heap2.heappop())

    def findMedian(self) -> float:
        if self.heap1.len() == self.heap2.len() + 1:
            return self.heap1.top()
        elif self.heap2.len() == self.heap1.len() + 1:
            return self.heap2.top()
        else:
            return (self.heap1.top() + self.heap2.top()) / 2


# Your MedianFinder object will be instantiated and called as such:
# obj = MedianFinder()
# obj.addNum(num)
# param_2 = obj.findMedian()