滑动窗口中位数

标签: 数组 哈希表 滑动窗口 堆(优先队列)

难度: Hard

中位数是有序序列最中间的那个数。如果序列的长度是偶数,则没有最中间的数;此时中位数是最中间的两个数的平均数。

例如:

  • [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 以内的答案将被视作正确答案。

Submission

运行时间: 145 ms

内存: 30.0 MB

import heapq
from collections import defaultdict
class Solution:
    def medianSlidingWindow(self, nums: List[int], k: int) -> List[float]:
      mp = defaultdict(int)
      bigger = []
      smaller = []
      balance = 0
      res = []
      for i,num in enumerate(nums):
        if i>=k:
          mp[nums[i-k]]+=1
          if nums[i-k]<=-smaller[0]:
            balance-=1
          else:
            balance+=1
        if not smaller or -smaller[0]>=num:
          heapq.heappush(smaller, -num)
          if len(smaller)+balance>1+len(bigger):
            heapq.heappush(bigger, -heapq.heappop(smaller))
        else:
          heapq.heappush(bigger, num)
          if len(smaller)+balance<len(bigger):
            heapq.heappush(smaller, -heapq.heappop(bigger))
        if i<k-1:
          continue
        while smaller and mp[-smaller[0]]:
          mp[-smaller[0]]-=1
          heapq.heappop(smaller)
          balance+=1
        while bigger and mp[bigger[0]]:
          mp[bigger[0]]-=1
          heapq.heappop(bigger)
          balance-=1
        if len(smaller)+balance == len(bigger):
          res.append((bigger[0]-smaller[0])/2)
        else:
          res.append(-smaller[0])
      return res

Explain

本题使用两个堆(大顶堆和小顶堆)来维护滑动窗口内的元素,并通过平衡两个堆的大小来计算中位数。具体思路如下: 1. 初始化大顶堆smaller和小顶堆bigger,用于存储滑动窗口内的元素。 2. 遍历数组nums,对于每个元素num: - 如果当前下标i>=k,说明滑动窗口向右移动,需要移除窗口最左边的元素nums[i-k],更新balance值。 - 根据num与堆顶元素的大小关系,将num插入到合适的堆中,并通过调整两个堆的大小来保持平衡。 - 如果当前下标i<k-1,说明还未形成完整的滑动窗口,继续遍历下一个元素。 - 移除堆顶失效的元素(即在窗口外的元素),并更新balance值。 - 根据两个堆的大小关系,计算当前窗口的中位数,并将其加入结果数组res中。 3. 返回结果数组res。

时间复杂度: O(nlogk)

空间复杂度: O(n)

```python
import heapq
from collections import defaultdict

class Solution:
    def medianSlidingWindow(self, nums: List[int], k: int) -> List[float]:
        mp = defaultdict(int)  # 哈希表,用于记录元素出现次数
        bigger = []  # 小顶堆,存储较大的一半元素
        smaller = []  # 大顶堆,存储较小的一半元素
        balance = 0  # 平衡因子,用于维护两个堆的大小关系
        res = []  # 结果数组,存储每个窗口的中位数
        
        for i, num in enumerate(nums):
            if i >= k:
                # 滑动窗口向右移动,移除最左边的元素
                mp[nums[i-k]] += 1
                if nums[i-k] <= -smaller[0]:
                    balance -= 1
                else:
                    balance += 1
            
            if not smaller or -smaller[0] >= num:
                # 将当前元素插入到smaller中
                heapq.heappush(smaller, -num)
                if len(smaller) + balance > 1 + len(bigger):
                    # 调整两个堆的大小,保持平衡
                    heapq.heappush(bigger, -heapq.heappop(smaller))
            else:
                # 将当前元素插入到bigger中
                heapq.heappush(bigger, num)
                if len(smaller) + balance < len(bigger):
                    # 调整两个堆的大小,保持平衡
                    heapq.heappush(smaller, -heapq.heappop(bigger))
            
            if i < k - 1:
                # 还未形成完整的滑动窗口,继续遍历下一个元素
                continue
            
            while smaller and mp[-smaller[0]]:
                # 移除smaller中失效的元素
                mp[-smaller[0]] -= 1
                heapq.heappop(smaller)
                balance += 1
            
            while bigger and mp[bigger[0]]:
                # 移除bigger中失效的元素
                mp[bigger[0]] -= 1
                heapq.heappop(bigger)
                balance -= 1
            
            if len(smaller) + balance == len(bigger):
                # 两个堆大小相等,中位数为两个堆顶元素的平均值
                res.append((bigger[0] - smaller[0]) / 2)
            else:
                # 两个堆大小不相等,中位数为smaller的堆顶元素
                res.append(-smaller[0])
        
        return res
```

Explore

在计算滑动窗口的中位数时,需要能够快速获取到窗口内的中间值。使用两个堆,一个大顶堆来存储窗口内较小的一半元素,一个小顶堆来存储窗口内较大的一半元素,可以有效地管理和平衡数据,使得能够快速从两个堆顶获取中位数。单个堆无法同时有效管理两端的数据,因此无法直接提供中位数的快速访问,尤其是在元素持续进出的动态窗口中。

在滑动窗口算法中,我们维护两个堆的元素数量之差不超过1。当插入新元素时,根据其大小与当前的堆顶元素比较决定插入哪个堆。如果新元素较大,则插入小顶堆;如果较小或等于,则插入大顶堆。插入后,如果发现一个堆的元素数量超过另一个堆的元素数量超过1,将数量多的堆的堆顶元素移动到另一个堆中,以恢复平衡。这样做保证了中位数始终位于两个堆的堆顶,可以快速获取。

为了高效处理失效的元素,我们使用一个哈希表记录每个元素在窗口外的出现次数。当窗口滑动导致元素移出时,该元素的计数增加。在从堆中取出堆顶元素时,检查该元素是否在哈希表中有记录(即是否失效)。如果失效,则继续弹出堆顶元素,直到找到一个有效的堆顶元素。这种方法避免了对堆进行全面重构,大大提高了处理失效元素的效率。