数组中滑动窗口的中位数

介绍

在一个数组中，选取长度为k的滑动窗口，从左往右滑动，每次滑动1个位置。求出每个滑动窗口的中位数。本题是数组中的滑动窗口的进阶版。

解法1：暴力枚举

对于每个滑动窗口，我们都可以利用快排的思想来求中位数。时间复杂度为O(klog(k))，总时间复杂度为O(nk*log(k))。代码如下：

class Solution:
    def medianSlidingWindow(self, nums: List[int], k: int) -> List[float]:
        def find_median(l):
            l.sort()
            if k % 2 == 1:
                return float(l[k//2])
            else:
                return (l[k//2-1]+l[k//2])/2
            
        res = []
        for i in range(len(nums)-k+1):
            res.append(find_median(nums[i:i+k]))
        return res

解法2：Heap

在解决此问题之前，我们先来看一个小根堆和大根堆的问题。

小根堆

小根堆的性质为：堆顶元素最小。当我们需要寻找一个序列中的最小值时，可以将序列中的元素全部放入小根堆中，然后不断弹出堆顶元素，直到找到最小值。时间复杂度为O(n*log(n))。

import heapq

lst = [2, 5, 1, 6, 8, 9]
heapq.heapify(lst)  # 转换为小根堆
print(heapq.heappop(lst))  # 输出1

大根堆

大根堆的性质为：堆顶元素最大。当我们需要寻找一个序列中的最大值时，可以将序列中的元素全部放入大根堆中，然后不断弹出堆顶元素，直到找到最大值。时间复杂度为O(n*log(n))。

import heapq

lst = [2, 5, 1, 6, 8, 9]
lst = [-x for x in lst]  # 将每个元素取负数
heapq.heapify(lst)  # 转换为大根堆
print(-heapq.heappop(lst))  # 输出9

比较器

Python的heapq模块是对堆的一种封装，其默认是小根堆。如果需要使用大根堆，我们可以先将每个元素取负数，然后再进行操作。

在Python中，还有一个非常重要的概念是比较器。比较器是指两个对象之间的比较关系。在Python中，可以用lambda表达式来定义一个比较器。

import heapq

lst = [(3, 'a'), (1, 'c'), (2, 'b')]
heapq.heapify(lst)  # 默认是利用元组的第一个元素排序，即先按数字从小到大排序
print(heapq.heappop(lst))  # 输出(1, 'c')
lst = [(3, 'a'), (1, 'c'), (2, 'b')]
heapq.heapify(lst, key=lambda x: -x[0])  # 通过比较器改造成大根堆
print(heapq.heappop(lst))  # 输出(3, 'a')

解题思路

根据上面所述，我们可以利用一个大根堆和一个小根堆，让它们分别负责前一半和后一半的元素，并保证大根堆的堆顶元素小于等于小根堆的堆顶元素。这样，我们就能在O(1)的时间内计算中位数了。

当我们把一个新的元素加入到滑动窗口中时，首先要判断它是在前一半还是后一半。如果前一半的元素数量小于等于后一半，就把新元素加入大根堆；否则加入小根堆。然后我们就需要对两个堆进行调整，使它们保持堆的性质。

当一个元素离开滑动窗口时，同样需要对堆进行调整。具体而言，如果它在大根堆，就从大根堆中删除它；如果它在小根堆，就从小根堆中删除它。然后，我们就需要对两个堆进行调整，使它们保持堆的性质。

具体实现可以参考代码：

class Solution:
    def medianSlidingWindow(self, nums: List[int], k: int) -> List[float]:
        if k % 2 == 0:
            is_even = True
        else:
            is_even = False

        max_heap = []  # 大根堆
        min_heap = []  # 小根堆

        # 初始化
        for i in range(0, k):
            heapq.heappush(max_heap, -nums[i])
        for i in range(0, k//2):
            val = -heapq.heappop(max_heap)
            heapq.heappush(min_heap, val)

        res = []

        # 滑动窗口
        for i in range(k, len(nums)):
            # 加入
            val = nums[i]
            if val <= -max_heap[0]:
                heapq.heappush(max_heap, -val)
            else:
                heapq.heappush(min_heap, val)

            # 调整
            if len(max_heap) < len(min_heap):
                heapq.heappush(max_heap, -heapq.heappop(min_heap))
            if len(max_heap) > len(min_heap) + 1:
                heapq.heappush(min_heap, -heapq.heappop(max_heap))

            # 移除
            if nums[i-k] <= -max_heap[0]:
                max_heap.remove(-nums[i-k])
                heapq.heapify(max_heap)
            else:
                min_heap.remove(nums[i-k])
                heapq.heapify(min_heap)

            # 计算中位数
            if is_even:
                median = (-max_heap[0] + min_heap[0]) / 2
            else:
                median = -max_heap[0]
            res.append(median)

        return res

总时间复杂度为O(n*log(k))。