Leetcode - 295. Find Median from Data Stream [H]

문제

Median 값은 다음을 의미한다.

• For example, for arr = [2,3,4], the median is 3.
• For example, for arr = [2,3], the median is (2 + 3) / 2 = 2.5.

값을 추가하고(addNum(int num)) Median값을 찾는(findMedian()) 함수를 구현하라.

Input
["MedianFinder", "addNum", "addNum", "findMedian", "addNum", "findMedian"]
[[], [1], [2], [], [3], []]
Output
[null, null, null, 1.5, null, 2.0]

Explanation
MedianFinder medianFinder = new MedianFinder();
medianFinder.addNum(1);    // arr = [1]
medianFinder.addNum(2);    // arr = [1, 2]
medianFinder.findMedian(); // return 1.5 (i.e., (1 + 2) / 2)
medianFinder.addNum(3);    // arr[1, 2, 3]
medianFinder.findMedian(); // return 2.0

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

해결 아이디어

• 두개의 heap을 사용하라는 힌트를 받음.
• 처음에 k번째로 큰 값을 찾는 k크기의 min heap 기법을 사용하려고했으나 실패.
• 1 2 3 (4) max heap 과 (6) 7 8 9 min heap 이렇게 두개의 heap을 사용해 중간 값을 찾는다. addNum()이 호출될때마다 이 두 heap을 유지하는 방법이 잘 떠오르지 않아 해설 참고함.
• 이 두개의 heap에 값을 추가하고, 두 heap의 크기 차이가 1이하를 유지 하도록 밸런스를 어떻게 맞추는지가 핵심. 코드를 잘 살펴보고 이해하기. 값을 하나 추가하고 그 뒤에 밸런스를 맞추면됨.
• 좋은 문제라고 생각함!

해결

class MedianFinder {
    //      max   min
    //  1 2 (3)   (4) 5 6
    priority_queue<int, vector<int>, greater<int>> minh;
    priority_queue<int, vector<int>, less<int>> maxh;
public:
    MedianFinder() {
        
    }
    
    void addNum(int num) {
        // 1. insert
        if (maxh.empty() || num > maxh.top()) {
            minh.push(num);
        } else {
            maxh.push(num);
        }
        // 2. re-balancing
        if (maxh.size() > minh.size() + 1) {
            minh.push(maxh.top());
            maxh.pop();
        } else if (maxh.size() + 1 < minh.size()) {
            maxh.push(minh.top());
            minh.pop();
        }
    }
    
    double findMedian() {
        if (maxh.size() > minh.size())    
            return (double)maxh.top();
        else if (maxh.size() < minh.size())    
            return (double)minh.top();
        else
            return (double)((maxh.top() + minh.top()) * 0.5);
    }
};

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