数组中滑动窗口的中位数

问题描述

给定一个数组nums，和一个大小为k的窗口从数组的最左端滑动到最右端，你只能看到在窗口内的k个数字。每次滑动窗口向右移动一个位置。找到每次滑动后窗口内的中位数。

思路

对于一个有序数组，中位数可以很容易地确定。然而，由于我们每次需要获取k个元素的中位数，因此，每次都对整个窗口进行排序的方法显然不是最优的。

我们为了找到中位数，可以先对整个窗口进行排序，找到排序后位于中间位置的元素，即中位数。然而，插入一个新元素和删除一个旧元素都需要重新排序，时间复杂度为 O(klogk)。我们可以使用数据结构来优化这一过程。

维护两个堆max_heap和min_heap，其中，max_heap中存储窗口中最大的一半元素，min_heap中存储窗口中最小的一半元素。我们规定元素全部存储在max_heap中，如果max_heap中元素个数多于min_heap，将max_heap的堆顶元素pop到min_heap。相应的，如果min_heap中元素多，将min_heap的堆顶元素pop到max_heap。以此保证两个堆存储的元素个数尽量平均。

当k为偶数时，中位数为max_heap的堆顶与min_heap的堆顶的平均值。因为，max_heap的堆顶一定小于min_heap的堆顶，因此这是中位数的计算方法。当k为奇数时，中位数为max_heap的堆顶。

最后，我们只需要在向右移动窗口时，将堆的状态进行调整，然后求出中位数即可。

代码实现

import heapq
from typing import List

class MedianFinder:
    def __init__(self):
        """
        initialize your data structure here.
        """
        self.max_heap = []
        self.min_heap = []

    def addNum(self, num: int) -> None:
        heapq.heappush(self.max_heap, -num)
        heapq.heappush(self.min_heap, -heapq.heappop(self.max_heap))
        if len(self.min_heap) > len(self.max_heap):
            heapq.heappush(self.max_heap, -heapq.heappop(self.min_heap))

    def findMedian(self) -> float:
        if len(self.max_heap) == len(self.min_heap):
            return (-self.max_heap[0] + self.min_heap[0])/2
        else:
            return -self.max_heap[0]

def medianSlidingWindow(nums: List[int], k: int) -> List[float]:
    res = []
    win = MedianFinder()
    for i, x in enumerate(nums):
        win.addNum(x)
        if i >= k:
            win.max_heap.remove(-nums[i-k])
            heapq.heapify(win.max_heap)
        if i >= k-1:
            res.append(win.findMedian())
    return res

复杂度分析

时间复杂度：O(NlogK)，其中N为数组的长度，k为窗口大小。因为每个元素最多只会被加入或移出一个堆，而插入和删除一个堆中元素的时间复杂度均为O(logK)。因此，总时间复杂度为O(NlogK)。

空间复杂度：O(K)， 取决于我们最终返回的中位数列表的大小是多少。我们只需要为两个堆分别维护 k//2 个数字。