技術標籤：每日一題java演算法leetcode資料結構佇列

480. 滑動視窗中位數

中位數是有序序列最中間的那個數。如果序列的大小是偶數，則沒有最中間的數；此時中位數是最中間的兩個數的平均數。

例如：

• [2,3,4]，中位數是 3
• [2,3]，中位數是 (2 + 3) / 2 = 2.5

給你一個數組 nums，有一個大小為 k 的視窗從最左端滑動到最右端。視窗中有 k 個數，每次視窗向右移動 1 位。你的任務是找出每次視窗移動後得到的新視窗中元素的中位數，並輸出由它們組成的陣列。

示例：

給出 nums = [1,3,-1,-3,5,3,6,7]，以及 k = 3。

視窗位置                      中位數
--
 
-------------               -----
[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
 因此，返回該滑動視窗的中位數陣列 [1,-1,-1,3,5,6]。

提示：

• 你可以假設 k 始終有效，即：k 始終小於輸入的非空陣列的元素個數。
• 與真實值誤差在 10 ^-5 以內的答案將被視作正確答案。

方法一：優先佇列

解題思路

此題的暴力解法是很容易想出來的：維護一個長度為 k 的陣列，左出右進，每次排序後取中位數即可，但這樣的時間複雜度是 O(klogk)。在此基礎上，思考優化解法。

我們知道 優先佇列 中的元素是有序的，但是優先佇列並不支援按索引取值。為此，可以維護兩個優先佇列 smallHeap 和 largeHeap，各自存一半的資料，如果 k 為奇數 那麼smallHeap 多放一個。注意：

1. smallHeap 是維護較小的那一半資料，堆頂是最大值，所以 smallHeap 是一個「大頂堆」
2. largeHeap 是維護較大的那一半資料，堆頂是最小值，所以 largeHeap
   是一個「小頂堆」
3. k 是奇數，中位數就是 smallHeap.peek()；反之，中位數就是 (smallHeap.peek() + largeHeap.peek()) / 2

儲存的問題搞定了，還有刪除。優先佇列只能刪除堆頂元素。為此可以額外使用一個 Map 維護需要刪除的元素，key 為待刪除的數字，value 為此數字需要刪除的次數。當堆頂元素在 Map 中時，說明此時需要刪除堆頂的元素。

此外，我們還需要保持兩個佇列的平衡，即 smallHeap.size() 要麼等於 largeHeap.size() ，要麼等於 largeHeap.size() + 1。因為兩個佇列都可能存在 “待刪除元素”，所以不能直接使用 size，需要手動維護兩個佇列中 “真實元素” 的數量。

參考程式碼

• 程式碼中封裝了 DualHeap 類，其中 #insert(int x)、#remove(int x)、#getMedian() 分別代表 插入元素、刪除元素、獲取中位數。
• 每次插入元素和刪除元素後，都需要呼叫 #balance() 方法平衡兩個佇列以及#checkTop() 方法檢查堆頂元素是否可以刪除。

public double[] medianSlidingWindow(int[] nums, int k) {
    int n = nums.length;
    double[] ans = new double[n - k + 1];
    DualHeap dualHeap = new DualHeap((k & 1) == 1);
    for (int i = 0; i < n; i++) {
        dualHeap.insert(nums[i]);
        if (i < k - 1) {
            continue;
        }
        ans[i - k + 1] = dualHeap.getMedian();
        dualHeap.remove(nums[i - k + 1]);
    }
    return ans;
}

@SuppressWarnings("all")
class DualHeap {
	// k 是否為奇數
    private boolean odd;
    private PriorityQueue<Integer> smallHeap;
    private PriorityQueue<Integer> largeHeap;
    private Map<Integer, Integer> lazyMap;
    // 兩個佇列中真實的元素數量（去掉待刪除的）
    int smallSize, largeSize;

    public DualHeap(boolean odd) {
        this.odd = odd;
        this.smallHeap = new PriorityQueue<>(Comparator.reverseOrder());
        this.largeHeap = new PriorityQueue<>(Comparator.naturalOrder());
        this.lazyMap = new HashMap<>();
        smallSize = 0; largeSize = 0;
    }

    public void insert(int val) {
        if (smallHeap.isEmpty() || val <= smallHeap.peek()) {
            smallSize++;
            smallHeap.offer(val);
        } else {
            largeSize++;
            largeHeap.offer(val);
        }
        balance();
        checkTop();
    }

    public void remove(int val) {
        if (val <= smallHeap.peek()) {
            smallSize--;
        } else {
            largeSize--;
        }
        lazyMap.put(val, lazyMap.getOrDefault(val, 0) + 1);
        balance();
        checkTop();
    }

    public double getMedian() {
        return odd ? smallHeap.peek() : ((double) smallHeap.peek() + largeHeap.peek()) / 2;
    }

    // 執行插入和刪除操作後，需要維持兩個佇列的平衡
    private void balance() {
        if (smallSize > largeSize + 1) {
            smallSize--;
            largeSize++;
            largeHeap.offer(smallHeap.poll());
        } else if (smallSize < largeSize) {
            smallSize++;
            largeSize--;
            smallHeap.offer(largeHeap.poll());
        }
    }

    // 如果堆頂元素是待刪除元素則刪除
    private void checkTop() {
        while (!smallHeap.isEmpty() && lazyMap.containsKey(smallHeap.peek())) {
            int x = smallHeap.poll();
            int count = lazyMap.get(x);
            if (count == 1) {
                lazyMap.remove(x);
            } else {
                lazyMap.put(x, count - 1);
            }
        }
        while (!largeHeap.isEmpty() && lazyMap.containsKey(largeHeap.peek())) {
            int x = largeHeap.poll();
            int count = lazyMap.get(x);
            if (count == 1) {
                lazyMap.remove(x);
            } else {
                lazyMap.put(x, count - 1);
            }
        }
    }
}

執行結果