295. Find Median from Data Stream

  • Two priority queues:

    • A max-heap lo to store the smaller half of the numbers
    • A min-heap hi to store the larger half of the numbers
  • The max-heap lo is allowed to store, at worst, one more element more than the min-heap hi. Hence if we have processed k elements:

    1. If , then lo is allowed to hold elements, while hi can hold nn elements.
    2. If ,then both heaps are balanced and hold nn elements each.

This gives us the nice property that when the heaps are perfectly balanced, the median can be derived from the tops of both heaps. Otherwise, the top of the max-heap lo holds the legitimate median.

  • Adding a number num:

    • Add num to max-heap lo. Since lo received a new element, we must do a balancing step for hi. So remove the largest element from lo and offer it to hi.
    • The min-heap hi might end holding more elements than the max-heap lo, after the previous operation. We fix that by removing the smallest element from hi and offering it to lo.

The above step ensures that we do not disturb the nice little size property we just mentioned

class MedianFinder {
public:

    // Adds a number into the data structure.
    void addNum(int num) {
        small.push(num);
        large.push(-small.top());
        small.pop();
        if (small.size() < large.size()) {
            small.push(-large.top());
            large.pop();
        }
    }

    // Returns the median of current data stream
    double findMedian() {
        return small.size() > large.size() ? small.top() : 0.5 *(small.top() - large.top());
    }

private:
    priority_queue<long> small, large;
};

results matching ""

    No results matching ""