Find Median from Data Stream

Design a data structure that supports the following two operations:

1. void addNum(int num): Adds an integer num to the data structure.
2. double findMedian(): Returns the median of all elements so far. The median is the middle value of a sorted dataset. If the number of elements is even, the median is the average of the two middle values.

Sample Input/Output:

Let’s consider a series of operations on the data stream:

addNum(1)
addNum(2)
findMedian() -> Output: 1.5
addNum(3)
findMedian() -> Output: 2

Explanation:
After adding the numbers 1 and 2, the median is (1 + 2) / 2 = 1.5 (average of the two middle values). After adding 3, the median becomes 2 (the middle value).

Time and Space Complexity:

The time complexity for adding a number is O(log n), where n is the total number of elements in the data stream. The time complexity for finding the median is O(1).

The space complexity is O(n) to store the elements in a data structure.

C++ Solution:

#include <iostream>
#include <queue>

using namespace std;

class MedianFinder {
public:
    priority_queue<int> maxHeap; // To store the smaller half of the data stream
    priority_queue<int, vector<int>, greater<int>> minHeap; // To store the larger half of the data stream

    void addNum(int num) {
        if (maxHeap.empty() || num <= maxHeap.top()) {
            maxHeap.push(num);
        } else {
            minHeap.push(num);
        }

        // Balance the heaps to ensure the size difference is at most 1
        if (maxHeap.size() > minHeap.size() + 1) {
            minHeap.push(maxHeap.top());
            maxHeap.pop();
        } else if (minHeap.size() > maxHeap.size()) {
            maxHeap.push(minHeap.top());
            minHeap.pop();
        }
    }

    double findMedian() {
        if (maxHeap.size() == minHeap.size()) {
            return (maxHeap.top() + minHeap.top()) / 2.0;
        } else {
            return maxHeap.top();
        }
    }
};

int main() {
    MedianFinder mf;

    mf.addNum(1);
    mf.addNum(2);
    cout << "Median: " << mf.findMedian() << endl; // Output: 1.5
    mf.addNum(3);
    cout << "Median: " << mf.findMedian() << endl; // Output: 2

    return 0;
}

In this implementation, we use two heaps to store the elements of the data stream. The maxHeap is used to store the smaller half of the data stream, while the minHeap is used to store the larger half.

When adding a number, we first determine which heap to add it to, based on whether it is smaller or greater than the top elements of the heaps. We then balance the heaps to ensure the size difference between the two heaps is at most 1.

The findMedian function returns the median based on the sizes of the two heaps. If the sizes are equal, we return the average of the two top elements. Otherwise, we return the top element of the maxHeap.

The output for the provided test case will be:

Median: 1.5
Median: 2