如何使用 Python 堆实现滑动窗口问题？

2023-04-11 00:00:00 窗口 如何使用 滑动

滑动窗口问题是指给定一个数组和滑动窗口的大小，求出所有滑动窗口内的最大值或最小值等问题。使用 Python 中的 heapq 模块可以方便地实现堆数据结构，进而解决滑动窗口问题。

下面是一个求出所有滑动窗口的最大值的示例代码：

import heapq

def max_sliding_window(nums, k):
    if not nums or k < 1 or len(nums) < k:
        return []
    heap = [(-n, i) for i, n in enumerate(nums[:k])]
    heapq.heapify(heap)
    res = [-heap[0][0]]
    for i in range(k, len(nums)):
        heapq.heappush(heap, (-nums[i], i))
        while heap[0][1] <= i - k:  # 弹出过期元素
            heapq.heappop(heap)
        res.append(-heap[0][0])
    return res

nums = [1, 3, -1, -3, 5, 3, 6, 7]
k = 3
print(max_sliding_window(nums, k))  # 输出 [3, 3, 5, 5, 6, 7]

这个函数的主要思路是维护一个大小为 k 的大根堆，每次滑动窗口时将一个新元素加入堆中，然后弹出过期元素（即下标不在当前窗口内的堆顶元素），并将堆顶元素的值存储到结果数组中。

在本例中，我们用到了 Python 中的 heapq 模块，它提供了三个主要函数：

  • heapq.heappush(heap, item)：将 item 加入 heap 中；
  • heapq.heappop(heap)：弹出 heap 的堆顶元素；
  • heapq.heapify(heap)：将 heap 转化为堆结构。

需要注意的是，heapq 模块中的函数操作的是小根堆，如果要实现大根堆，可以将元素值取负数。

相关文章