Leetcode - 480. Sliding Window Median 풀이 [H]

문제

주어진 배열을 순회하면서 k 크기의 window size 의 모든 median값을 구하라.

Input: nums = [1,3,-1,-3,5,3,6,7], k = 3
Output: [1.00000,-1.00000,-1.00000,3.00000,5.00000,6.00000]
Explanation:
Window position                Median
---------------                -----
[1  3  -1] -3  5  3  6  7        1
 1 [3  -1  -3] 5  3  6  7       -1
 1  3 [-1  -3  5] 3  6  7       -1
 1  3  -1 [-3  5  3] 6  7        3
 1  3  -1  -3 [5  3  6] 7        5
 1  3  -1  -3  5 [3  6  7]       6

https://leetcode.com/problems/sliding-window-median/

해결 아이디어

• sliding window 문제는 매번 윈도우 크기 k를 모두 계산할 필요가 없이. k범위의 이전값을 제거, 이후값을 추가하는 O(1)연산만 추가로 하면 되기에, 배열크기가 n일때 시간복잡도는 총 O(n)이 된다는 장점이 있음.

• 이 문제에서도 min heap / max heap두개를 만들고 k 범위 이전값을 heap에서 제거, 이후값을 heap에 추가한 뒤 heap을 re-balance하여 유지하면 된다.

• Leetcode - 295. Find Median from Data Stream [H] 문제와 거의 동일한 문제. 이 문제가 약간의 심화버전임. 같이 참고할것.

• 문제는 priority_queue는 top()이외의 값을 제거할 수가 없다는 점. del_num()동작을 할때, heap내에 있는 특정 값을 제거해야한다. 따라서 priority_queue 대신 multiset을 사용한다. 특정 element를 찾는 동작이 O(logn) , 제거 동작도 O(logn)

• multiset은 rbtree로 정렬되어있다. 따라서 첫번째 혹은 마지막 iterator가 가장 크거나 작은값을 나타낸다. 이 성질을 이용해 heap으로 사용하는 것.

      maxh     minh
     1 2 (3)   (4) 5 6
     |    |
     |    maxh.rbegin() or --maxh.end()
     maxh.begin()

• maxh.erase() 의 인자로 iterator를 넣을 수 있음.
• rbegin() 과 rend()는 std::reverse_iterator 타입 이다. 역방향의 시작과 끝의 reverse_iterator 를 리턴한다. 하지만 그것을 erase()의 인자로는 전달 할수 없다. iterator 타입만 전달받을 수 있기 때문. 따라서 maxh.erase(--maxh.end()); 와 같이 iterator 활용.

해결

class Solution {
    //  maxh     minh
    // 1 2 (3)   (4) 5 6
    // |    |
    // |    maxh.rbegin() or --maxh.end()
    // maxh.begin()
    multiset<int, less<int>> maxh; // inc order
    multiset<int, less<int>> minh;
public:
    void add_num(int val) {
        // 1. insert
        if (maxh.empty() || val > *maxh.rbegin())   
            minh.insert(val);
        else
            maxh.insert(val);
        
        // 2. re-balance
        if (maxh.size() > minh.size() + 1) {
            minh.insert(*(--maxh.end()));
            maxh.erase(--maxh.end()); // ~~FIXME: modify to O(1)~~
        } else if (maxh.size() + 1 < minh.size()) {
            maxh.insert(*minh.begin());
            minh.erase(minh.begin());
        }
    }
    void del_num(int val) {
        auto max_it = maxh.find(val);  // heap에서 val값을 찾아 제거해야함.
        auto min_it = minh.find(val);
        // 1. delete
        if (max_it != maxh.end()) {
            maxh.erase(max_it);
        } else if (min_it != minh.end()) {
            minh.erase(min_it);
        }
        // 2. re-balance
        if (maxh.size() > minh.size() + 1) {
            minh.insert(*(--maxh.end()));
            maxh.erase(--maxh.end()); // ~~FIXME: modify to O(1)~~
        } else if (maxh.size() + 1 < minh.size()) {
            maxh.insert(*minh.begin());
            minh.erase(minh.begin());
        }
    }
    
    double get_median(void) {
        if (maxh.size() > minh.size())
            return (double)*maxh.rbegin();
        else if (maxh.size() < minh.size())
            return (double)*minh.begin();
        else
            return ((double)*maxh.rbegin() + (double)*minh.begin()) * 0.5;
    }
    
    vector<double> medianSlidingWindow(vector<int>& nums, int k) {
        vector<double> ret;
        int nsize = nums.size();
        
        for (int i = 0; i < k; i++)
            add_num(nums[i]);
        ret.push_back(get_median());
        
        for (int i = 1; i < nsize - k + 1; i++) {
            del_num(nums[i - 1]);
            add_num(nums[i + k - 1]);
            ret.push_back(get_median());
        }
        return ret;
    }
};