LeetCode 239 滑动窗口最大值- Python手撕最大堆

Posted 2023-02-07 LeonYi

手撕版

最大堆的完全实现， 堆中元素为二元组（num, idx），比较时用数值，赋值或交换时用整个元组。

class Heap:
    def __init__(self, arr, capacity):
        # 容量和大小
        self.size = len(arr)
        self.arr = [None] * capacity
        self.arr[0] = (10e5, 0)
        for i, num in enumerate(arr):
            self.arr[i+1] = num 

    def heapify(self, parent):
        x = self.arr[parent]
        while parent * 2 <= self.size:
            child = parent * 2
            if child != self.size and self.arr[child + 1][0] > self.arr[child][0]:
                child += 1
            if self.arr[child][0] > x[0]:
                self.arr[parent] = self.arr[child]  # 孩子节点值上移
                parent = child
            else:
                break
        self.arr[parent] = x

    def insert(self, item):
        self.size += 1
        child = self.size  # 空穴位置
        while item[0] > self.arr[child // 2][0]:
            parent = child // 2
            self.arr[child] = self.arr[parent]
            child = parent
        self.arr[child] = item

    def pop(self):
        max_item = self.arr[1]  # 取堆顶
        self.arr[1] = self.arr[self.size]  # 取堆末尾元素
        self.size -= 1
        self.heapify(1)
        # print(self.arr)
        return max_item

利用最大堆实现滑动窗口最大值。

创建一个初始大小为K的最大堆，然后向右移动逐个添加元素，同时根据元素索引判断堆顶元素是否在滑动窗口内，若不再则出堆直到在范围内。

class Solution:
    def maxSlidingWindow(self, nums: List[int], k: int) -> List[int]:
        if k == 1 or len(nums) == 1:
            return nums
        else:
            res = []
            n = len(nums)
            # 构造元组形式输入
            arr = [(nums[j], j + 1) for j in range(k)]

            h = Heap(arr, n+1)
            for j in range(k // 2, 0, -1):
                h.heapify(j)

            res.append(h.arr[1][0])

            for i in range(k, n):
                h.insert((nums[i], i + 1))
                # i - k + 1为元素索引从1开始的堆
                while h.arr[1][1] <= (i - k + 1):
                    h.pop()
                res.append(h.arr[1][0])

            return res

不看之前的代码，完全靠手撕最大堆有点费劲。除了上述实现，其实还可以每次向右滑动时，先删除之前K个元素的第一个，然后继续添加元素，再取堆顶。不过手撕起来比较复杂。

调包版

如果面试官允许，可以调包...
heapq的heapify、heappush和heappop

class Solution:
    def maxSlidingWindow(self, nums: List[int], k: int) -> List[int]:
        if k == 1 or len(nums) == 1:
            return nums
        else:
            import heapq
            arr = [(-nums[j], j + 1) for j in range(k)]
            heapq.heapify(arr)
            res = [-arr[0][0]]

            for i in range(k, len(nums)):
                heapq.heappush(arr, (-nums[i], i + 1))
                # i - k + 1为元素索引从1开始的堆
                while arr[0][1] <= (i - k + 1):
                    heapq.heappop(arr)
                res.append(-arr[0][0])

            return res