[LeetCode] 295. Find Median from Data Stream

Median is the middle value in an ordered integer list. If the size of the list is even, there is no middle value. So the median is the mean of the two middle value.
Examples: 
[2,3,4] , the median is 3
[2,3], the median is (2 + 3) / 2 = 2.5
Design a data structure that supports the following two operations:
  • void addNum(int num) - Add a integer number from the data stream to the data structure.
  • double findMedian() - Return the median of all elements so far.
For example:
addNum(1)
addNum(2)
findMedian() -> 1.5
addNum(3) 
findMedian() -> 2

Thought process:

Use a max heap to maintain the smaller half of the numbers, and a min heap to maintain the larger half of the numbers. Designate one of the heap's top to be the median, or smaller of the middle values if total number is even. To maintain two heaps' balance, we need to make sure that one of the heap's size is always equal to or 1 larger than the other's. We can designate the heap that keeps track of the smaller half to be the larger heap. Once a number is added, it goes to the "small" heap first. Then, "small" heap should pass its largest element to "large", in order to maintain our partition. Lastly, if small's size is less than large's, pass large's top element (min) to small.

Note that priority queue is max heap in default in C++. To create a min heap, pass the greater<> function from the <functional> library to the heap constructor.

Also, if there is an even number of numbers, adding two middle numbers may cause integer overflow. I should use long to store the numbers.

Solution:
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
class MedianFinder {
private:
    priority_queue< long > small;
    priority_queue< long, vector< long >, greater< long > > large;
    
public:
    /** initialize your data structure here. */
    MedianFinder() {}
    
    void addNum(int num) {
        small.push(num);
        large.push(small.top());
        small.pop();
        if (small.size() < large.size()) {
            small.push(large.top());
            large.pop();
        }
    }
    
    double findMedian() {
        if (small.size() > large.size()) {
            return small.top();
        } else {
            return (small.top() + large.top()) / 2.0;
        }
    }
};

/**
 * Your MedianFinder object will be instantiated and called as such:
 * MedianFinder obj = new MedianFinder();
 * obj.addNum(num);
 * double param_2 = obj.findMedian();
 */

Time complexity:

addNum() takes O(logn). findMedian() takes O(1).

The problem can also be solved using binary search trees. Set in C++ is implemented using balanced BST. Since set does not permit duplicate elements, we will use multiset, which during insertion always adds the equal element to the right. With these nice properties, we can use an iterator to keep track of the median. We need to make sure that if there is an even number of numbers, the iterator always points to the smaller of the 2 middle numbers (or the larger one, depending on your choice).

Solution:
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
class MedianFinder {
private:
    multiset< long > set;
    multiset< long >::iterator mid;
    
public:
    /** initialize your data structure here. */
    MedianFinder() {}
    
    void addNum(int num) {
        if (set.size() == 0) {
            mid = set.insert(num);
        } else {
            set.insert(num);
            
            if ((set.size() & 1) && num >= *mid) {
                mid++;
            } else if (!(set.size() & 1) && num < *mid) {
                mid--;
            }
        }
    }
    
    double findMedian() {
        if (set.size() & 1) {
            return *mid;
        } else {
            return (*mid + *next(mid)) / 2.0;
        }
    }
};

/**
 * Your MedianFinder object will be instantiated and called as such:
 * MedianFinder obj = new MedianFinder();
 * obj.addNum(num);
 * double param_2 = obj.findMedian();
 */

Time complexity:

addNum() takes O(logn). findMedian() takes O(1).

Comments

Popular posts from this blog

[LeetCode] 714. Best Time to Buy and Sell Stock with Transaction Fee

[LeetCode] 269. Alien Dictionary

[LeetCode] 631. Design Excel Sum Formula