480. Sliding Window Median

題目網址：https://leetcode.cn/problems/sliding-window-median

題意：中位數是指有序 array 裡中間的數。若 array 長度偶數, 則中位數是中間兩數的平均值。

arr = [2,3,4], 則中位數 = 3

arr = [2,3], 則中位數 = (2+3) / 2 = 2.5

給一整數 array nums 和一整數 k, 有一長度為 k 的 sliding window 從 nums 最左端滑到最右端, 每次向右移動一位, 求每次移動 sliding window 後得到的中位數, 並返回由這些中位數所組成的 array。

Solution：

想法：利用 Heap, 類似 295. Find Median from Data Stream, 使用兩個 heap + 延遲刪除, 我們定義平衡：

滿足 smaller.size() - larger.size() ≤ 1

也就是說 smaller 的元素個數要麼比 larger 多一, 要麼相同

這樣設計的好處在於：

當 k 為偶數時, 各自的 top 之平均值即為中位數

當 k 為奇數時, 則 smaller.top() 即為中位數

每次移動 window 的操作：

在每次移動 sliding window 前都為平衡的, 也就是 smaller.size() - larger.size() ≤ 1

設置 balance = 0, 用來記錄該次移動 window 中, smaller 和 larger 元素個數的差值

刪除 window 左側的元素

由於 heap 沒辦法刪除指定的元素 val, 因此先欠下這個帳, 等到 val 出現在 heap.top() 時, 再進行刪除。

delay 紀錄 val 欠賬的次數, 因此 ++delay[val]

新增 window 右側的元素

若 nums[i] <= smaller.top(), 則把 nums[i] 放到 smalller 中, 然後 ++balance

否則, 把 nums[i] 放到 larger 中, 然後 -balance

經過上述刪除/新增的操作, 此時 balance 為 2、0、2 這三種情況之一：

2：本次移動中刪除的元素是在 smaller 中的, 且新增的元素是 larger

0：

本次移動中刪除的元素是在 smaller 中的, 且新增的元素是在 smaller

本次移動中刪除的元素是在 larger 中的, 且新增的元素是在 larger

2：本次移動中刪除的元素是在 larger 中的, 且新增的元素是 smaller

我們需要進行調整, 使 balance = 0, 也就 smaller.size() - larger.size() ≤ 1:

若 balance== -2, 則把 larger.top() 給 pop 掉, 並將其 push 到 smaller 中

若 balance== 0, 則不需調整

若 balance == 2, 則把 smaller.top() 給 pop 掉, 並將其 push 到 larger 中

調整完後, 接著要計算中位數。但在計算中位數前, 必須先把欠的帳給還了, 避免那些應該刪除的數影響到中位數的計算

分別檢查兩邊的 heap, 若 heap.top() 欠著債, 則將其 pop 掉, 直到 heap.top() 沒有欠債為止

若那些應刪除的數不為 heap.top() 該怎麼辦 ? 由於計算中位數的時候只與 heap.top() 有關, 至於那些不為 heap.top() 但欠著債的, 就讓他們欠著吧, 等他們成為 heap.top() 的時候再 pop 掉就行了

最後, 計算中位數

class Solution {
public:
    vector<double> medianSlidingWindow(vector<int>& nums, int k) {
        const int n = nums.size();
        vector<double> res;

        // 將 k 個元素 push 到 smaller 中
        for (int i = 0; i < k; ++i) {
            smaller.emplace(nums[i]);
        }

        // 將 (k / 2) 個元素從 smaller 中 pop 掉, 並 push 到 larger 中
        // 若 k 為奇數, 則此時 smaller 的元素個數比 larger 多一
        for (int i = 0; i < k / 2; ++i) {
            larger.emplace(smaller.top());
            smaller.pop();
        }

        // 得到第一個中位數
        res.emplace_back(getMedian(k));

        // 得到剩下 (n - k) 個中位數, 總共 (n - k + 1) 個中位數
        for (int i = k; i < n; ++i) {
            int balance = 0; // 每次插入 nums[i] 前都為平衡的

            // 判斷從哪個 heap 中刪除 nums[i-k]
            remove(nums[i-k], balance);

            // 判斷 nums[i] 要加入到哪個 heap 中
            insert(nums[i], balance);

            // 在刪除/新增元素後, 要仍保持平衡
            makeBalance(balance);

            // 若 heap.top() 恰為要刪除的數, 則要在計算中位數前刪除掉
            prune(smaller);
            prune(larger);

            res.emplace_back(getMedian(k));
        }

        return res;
    }

private:
    unordered_map<int, int> delay; // 紀錄要刪除的數, 以及其次數
    priority_queue<int> smaller; // max heap
    priority_queue<int, vector<int>, greater<int>> larger; // min heap

    double getMedian(int& k){
        if (k % 2 == 1) {
            return smaller.top();
        }
        return smaller.top() * 0.5 + larger.top() * 0.5; // 避免 overflow 的寫法
    }

    void insert(int& val, int& balance){
        if (!smaller.empty() && val <= smaller.top()) {
            smaller.emplace(val);
            ++balance;
        } else {
            larger.emplace(val);
            --balance;
        }
    }

    void remove(int& val, int& balance){
        if (!smaller.empty() && val <= smaller.top()) {
            --balance;
        } else {
            ++balance;
        }

        ++delay[val]; // 在 delay 中紀錄 val
    }

    void makeBalance(int& balance){
        if (balance > 0) {
            larger.emplace(smaller.top());
            smaller.pop();
        } else if (balance < 0) {
            smaller.emplace(larger.top());
            larger.pop();
        }
    }

    template<typename T>
    void prune(T& heap){
        while (!heap.empty() && delay[heap.top()] > 0) {
            --delay[heap.top()];
            heap.pop();
        }
    }
};

time：$O(n \cdot log(n))$ ➔ 每次 insert / delete 皆須 $log(n)$, worse case：應被刪除的數皆不為 top, 導致 smaller + larger 為 n 個數
space：$O(n)$ ➔ smaller + larger 和 delay, 長度皆不超過 n