Leetcode 480 Solution

This article provides solution to leetcode question 480 (sliding-window-median)
Categories:
Tags:
Access this page by simply typing in "lcs 480" in your browser address bar if you have bunnylol configured.

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

Solution

class MedianFinder {
    map<int, int> num_cnt;
    map<int, int>::iterator mid_it;
    int left_size;
    int right_size;

public:
    MedianFinder()
    {
        left_size = 0;
        right_size = 0;
    }

    // Adds a number into the data structure.
    void addNum(int num) {
        if (num_cnt.empty())
        {
            num_cnt[num]++;
            mid_it = num_cnt.begin();
        }
        else
        {
            num_cnt[num]++;

            if (num < mid_it->first)
                left_size++;
            else if (num > mid_it->first)
                right_size++;

            if (left_size == mid_it->second + right_size)
            {
                right_size += mid_it->second;
                mid_it--;
                left_size -= mid_it->second;
            }
            else if (left_size + mid_it->second + 1 == right_size)
            {
                left_size += mid_it->second;
                mid_it++;
                right_size -= mid_it->second;
            }
        }
    }

    void delNum(int num) {
        num_cnt[num]--;

        if (num < mid_it->first)
            left_size--;
        else if (num > mid_it->first)
            right_size--;

        if (left_size == mid_it->second + right_size)
        {
            right_size += mid_it->second;
            mid_it--;
            left_size -= mid_it->second;
        }
        else if (left_size + mid_it->second + 1 == right_size)
        {
            left_size += mid_it->second;
            mid_it++;
            right_size -= mid_it->second;
        }

        if (num_cnt[num] == 0)
            num_cnt.erase(num);
    }

    // Returns the median of current data stream
    double findMedian() {
        if (left_size + mid_it->second == right_size)
        {
            auto next_it = mid_it;
            next_it++;

            return (mid_it->first + (int64_t)next_it->first) * 0.5;
        }
        else
            return mid_it->first;
    }
};

class Solution {
public:
    vector<double> medianSlidingWindow(vector<int>& nums, int k) {
        MedianFinder mf;
        vector<double> res;

        for (int i = 0; i < nums.size(); i++)
        {
            mf.addNum(nums[i]);

            if (i >= k - 1)
            {
                res.push_back(mf.findMedian());
                mf.delNum(nums[i - (k - 1)]);
            }
        }

        return res;
    }
};