LeetCode 专题暴击:最大最小heap!

发布于 2023-05-25  1271 次阅读


前言

在这个专题,我们将主要讨论 heap 的具体用法。

有很多类型题会使用到 heap, 比如 top k 问题,合并 k 个有序链表,合并 k 个有序数组等等。

为什么使用他们呢?因为他们可以帮助我们快速的找到最大值或者最小值。

所以,在我们看到问题的一些字眼时,比如 top k,第 k 大,第 k 小,我们就可以想到 heap。

什么是 heap

heap 是一种数据结构,它的特点是可以在 O(1)的时间复杂度内找到最大值或者最小值。

在 python 里,我们可以使用 heapq 来实现 heap。我们接下来会讲解 heap 的常用方法。

注意,heapq 是小顶堆,如果我们需要大顶堆,我们需要把第一个元素取负数,然后再插入 heap。

import heapq

# 初始化一个 heap
heap = []

# 往 heap 里插入一个元素,然后 heap 会自动调整,复杂度为 O(logn)
heapq.heappush(heap, 1)

# 从 heap 里弹出最小的元素,复杂度为 O(logn)
heapq.heappop(heap)

# 从 heap 里取出最小的元素,但是不弹出,复杂度为 O(1)
heap[0]

# 把 heap 的最小的元素替换成 newvalue,然后弹出最小的元素,然后 heap 会自动调整,复杂度为 O(logn)
heapq.heapreplace(heap, newvalue)

# 把 heap 先push一个 newitem,然后再弹出最小的元素,复杂度为 O(logn)
heapq.heappushpop(heap, newitem)

# heapreplace 和 heappushpop的区别是,当 newvalue < heap[0] 时,heapreplace会先弹出最小的元素,然后再插入 newvalue,而 heappushpop 会直接返回 newvalue,不会插入 newvalue。

# 取 heap 中最大的 n 个元素,复杂度为 O(nlogn)
heapq.nlargest(n, heap)

# 取 heap 中最小的 n 个元素,复杂度为 O(nlogn)
heapq.nsmallest(n, heap)

接下来我们看几道例题!

215. Kth Largest Element in an Array

这道题的描述是找到数组中第 k 大的元素。

test cases:

Input: [3,2,1,5,6,4] and k = 2
Output: 5

Input: [3,2,3,1,2,4,5,5,6] and k = 4
Output: 4

这道题的思路是,我们可以使用一个大小为 k 的 heap,然后遍历数组,如果当前元素比 heap 的第一个元素大,说明 heap 的第一个元素不是第 k 大的元素,我们可以把第一个元素弹出,然后把当前元素插入 heap。

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        heap = []
        for num in nums:
            if len(heap) < k:
                heapq.heappush(heap, num)
            else:
                if num > heap[0]:
                    heapq.heappushpop(heap, num)
                    # 这里也可以使用 heapq.heappushpop(heap, num),区别是 heapreplace 会先弹出最小的元素,然后再插入 num,而 heappushpop 会直接返回 num,不会插入 num。
        return heap[0]

当然,我们也可以使用 heapq.nlargest 来解决这道题。

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        return heapq.nlargest(k, nums)[-1]

也可以反向思维,使用最大堆来解决这道题。

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        heap = [-num for num in nums]
        heapq.heapify(heap)
        for _ in range(k):
            res = heapq.heappop(heap)
        return -res

复杂度分析:

  • 时间复杂度:O(nlogk),原因是我们需要遍历整个数组,然后每次插入都需要 logk 的时间复杂度。
  • 空间复杂度:O(k),原因是我们需要使用 heap 来存储数字。

378. Kth Smallest Element in a Sorted Matrix

这道题的描述是找到一个有序矩阵中第 k 小的元素。

test cases:

Input:
matrix = [
   [ 1,  5,  9],
   [10, 11, 13],
   [12, 13, 15]
],
k = 8

Output: 13

这道题的思路是,我们可以使用一个大小为 k 的 heap,然后遍历矩阵,把每个矩阵的第一个元素放入 heap,然后每次从 heap 中取出最小的元素,然后把这个元素所在矩阵的下一个元素放入 heap。

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        heap = []
        for i in range(len(matrix)):
            heapq.heappush(heap, (matrix[i][0], i, 0)) # 先把每个矩阵的第一个元素放入 heap, 同时记录这个元素所在矩阵的坐标
        res = 0
        for _ in range(k): # 然后从 heap 中取出最小的元素,然后把这个元素所在矩阵的下一个元素放入 heap
            res, i, j = heapq.heappop(heap) # 这里的 i 和 j 是矩阵的坐标
            if j + 1 < len(matrix[0]): # 如果 j + 1 < len(matrix[0]),说明这个元素不是所在矩阵的最后一个元素,我们可以把这个元素所在矩阵的下一个元素放入 heap
                heapq.heappush(heap, (matrix[i][j + 1], i, j + 1))
        return res

复杂度分析:

  • 时间复杂度:O(klogk),原因是我们需要遍历整个矩阵,然后每次插入都需要 logk 的时间复杂度。
  • 空间复杂度:O(k),原因是我们需要使用 heap 来存储数字。

347. Top K Frequent Elements

这道题的描述是找到数组中出现频率最高的 k 个元素。

test cases:

Input: nums = [1,1,1,2,2,3], k = 2
Output: [1,2]

Input: nums = [1], k = 1
Output: [1]

这道题的思路是,我们可以使用一个大小为 k 的 heap,然后遍历数组,把每个元素放入 hash table,然后把 hash table 中的元素放入 heap,然后每次从 heap 中取出最小的元素,然后把这个元素所在 hash table 的下一个元素放入 heap。

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        heap = []
        count = collections.Counter(nums)
        for num, freq in count.items():
            if len(heap) < k:
                heapq.heappush(heap, (freq, num))
            else:
                if freq > heap[0][0]:
                    heapq.heappushpop(heap, (freq, num))
        return [num for freq, num in heap]

也可以使用 heapq.nlargest 来解决这道题。

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        count = collections.Counter(nums)
        return heapq.nlargest(k, count.keys(), key=count.get)

当然也可以使用最大堆来解决这道题。

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        heap = []
        count = collections.Counter(nums)
        for num, freq in count.items():
            heapq.heappush(heap, (-freq, num))
        res = []
        for _ in range(k):
            res.append(heapq.heappop(heap)[1])
        return res

复杂度分析:

  • 时间复杂度:O(nlogk),原因是我们需要遍历整个数组,然后每次插入都需要 logk 的时间复杂度。
  • 空间复杂度:O(k),原因是我们需要使用 heap 来存储数字。

692. Top K Frequent Words

这道题的描述是找到数组中出现频率最高的 k 个元素。

test cases:

Input: ["i", "love", "leetcode", "i", "love", "coding"], k = 2
Output: ["i", "love"]

Input: ["the", "day", "is", "sunny", "the", "the", "the", "sunny", "is", "is"], k = 4
Output: ["the", "is", "sunny", "day"]

这道题的思路是,我们可以使用一个大小为 k 的 heap,然后遍历数组,把每个元素放入 hash table,然后把 hash table 中的元素放入 heap,然后每次从 heap 中取出最小的元素,然后把这个元素所在 hash table 的下一个元素放入 heap。

!!注意这道题要求返回的是字母顺序最小的 k 个元素而且是需要从大到小排序的。所以我们只能使用最大堆来解决这道题。

class Solution:
    def topKFrequent(self, words: List[str], k: int) -> List[str]:
        heap = []
        count = collections.Counter(words)
        for word, freq in count.items():
            heapq.heappush(heap, (-freq, word))
        res = []
        for _ in range(k):
            res.append(heapq.heappop(heap)[1])
        return res

复杂度分析:

  • 时间复杂度:O(nlogk),原因是我们需要遍历整个数组,然后每次插入都需要 logk 的时间复杂度。
  • 空间复杂度:O(k),原因是我们需要使用 heap 来存储数字。

253. meeting rooms ii

这道题的描述是给定一组会议的开始和结束时间,求最少需要多少个会议室。

test cases:

Input: [[0, 30],[5, 10],[15, 20]]
Output: 2

Input: [[7,10],[2,4]]
Output: 1

这道题的思路是,我们把 interval 按照开始时间排序,然后使用一个 heap 来存储结束时间,然后遍历数组,如果当前的开始时间大于 heap 的最小值,那么就把 heap 的最小值 pop 出来,然后把当前的结束时间放入 heap。能这么做的原因是,如果当前的开始时间大于 heap 的最小值,那么说明这个会议可以使用之前的会议室,所以我们把之前的会议室的结束时间 pop 出来,然后把当前的结束时间放入 heap。这是一种变向的贪心算法。

class Solution:
    def minMeetingRooms(self, intervals: List[List[int]]) -> int:
        if not intervals:
            return 0
        intervals.sort()
        heap = []
        heapq.heappush(heap, intervals[0][1]) # 把第一个会议的结束时间放入 heap
        for i in range(1, len(intervals)):
            if intervals[i][0] >= heap[0]: # 如果当前的开始时间大于 heap 的最小值,那么就把 heap 的最小值 pop 出来,然后把当前的结束时间放入 heap
                heapq.heappop(heap) # 仔细想一下,其实就相当于抵消了一个会议室
            heapq.heappush(heap, intervals[i][1]) # 把当前的结束时间放入 heap
        return len(heap)

复杂度分析:

  • 时间复杂度:O(nlogn),原因是我们需要遍历整个数组,然后每次插入都需要 logn 的时间复杂度。
  • 空间复杂度:O(n),原因是我们需要使用 heap 来存储数字。

373. Find K Pairs with Smallest Sums

这道题的描述是找到两个数组中和最小的 k 对数。

test cases:


Input: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
Output: [[1,2],[1,4],[1,6]]

Input: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
Output: [1,1],[1,1]

这道题的思路是,我们可以使用一个大小为 k 的 heap,然后遍历数组,把所有的组合放入 heap,然后每次从 heap 中取出最小的元素。

class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        heap = []
        for num1 in nums1:
            for num2 in nums2:
                heapq.heappush(heap, (num1 + num2, [num1, num2]))
        res = []
        for _ in range(min(k, len(heap))):
            res.append(heapq.heappop(heap)[1])
        return res

但是这么做会超时,原因是我们把所有的组合都放入了 heap,但是我们只需要最小的 k 个组合,所以我们可以使用一个大小为 k 的 heap,然后遍历数组,把每个数组的第一个元素放入 heap,然后每次从 heap 中取出最小的元素,然后把这个元素所在数组的下一个元素放入 heap。

class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        res = []
        visited = set()

        min_heap = [(nums1[0] + nums2[0], (0, 0))]

        visited.add((0, 0))
        count = 0

        while k > 0 and min_heap:
            val, (i, j) = heapq.heappop(min_heap)
            res.append([nums1[i], nums2[j]])

            if i + 1 < len(nums1) and (i + 1, j) not in visited:
                heapq.heappush(min_heap, (nums1[i + 1] + nums2[j], (i + 1, j)))
                visited.add((i + 1, j))

            if j + 1 < len(nums2) and (i, j + 1) not in visited:
                heapq.heappush(min_heap, (nums1[i] + nums2[j + 1], (i, j + 1)))
                visited.add((i, j + 1))

            k -= 1

        return res

复杂度分析:

  • 时间复杂度:O(klogk),原因是我们需要遍历整个数组,然后每次插入都需要 logk 的时间复杂度。
  • 空间复杂度:O(k),原因是我们需要使用 heap 来存储数字。

23. Merge k Sorted Lists

这道题的描述是合并 k 个有序链表。

test cases:

Input:
[
  1->4->5,
  1->3->4,
  2->6
]
Output: 1->1->2->3->4->4->5->6

这道题的思路是,我们可以使用一个大小为 k 的 heap,然后遍历链表,把每个链表的第一个元素放入 heap,然后每次从 heap 中取出最小的元素,然后把这个元素所在链表的下一个元素放入 heap。

class Solution:
    def mergeKLists(self, lists: List[ListNode]) -> ListNode:
        heap = []
        for i, node in enumerate(lists):
            if node:
                heapq.heappush(heap, (node.val, i))
        dummy = ListNode(0)
        curr = dummy
        while heap:
            val, i = heapq.heappop(heap)
            curr.next = lists[i]
            curr = curr.next
            lists[i] = lists[i].next
            if lists[i]:
                heapq.heappush(heap, (lists[i].val, i))
        return dummy.next

复杂度分析:

  • 时间复杂度:O(nlogk),原因是我们需要遍历整个链表,然后每次插入都需要 logk 的时间复杂度。
  • 空间复杂度:O(k),原因是我们需要使用 heap 来存储数字。

407. Trapping Rain Water II

这道题的描述是给定一个二维数组,数组中的数字代表高度,然后问这个二维数组能够存储多少水。

test cases:

Given the following 3x6 height map:
[
  [1,4,3,1,3,2],
  [3,2,1,3,2,4],
  [2,3,3,2,3,1]
]

Return 4.

这道题的思路是,我们可以使用一个大小为 k 的 heap,然后遍历矩阵的四条边,把每个边的元素放入 heap,然后每次从 heap 中取出最小的元素,然后把这个元素所在矩阵的下一个元素放入 heap。

class Solution:
    def trapRainWater(self, heightMap: List[List[int]]) -> int:
        if not heightMap or not heightMap[0]:
            return 0

        heap = []
        visited = set()

        m, n = len(heightMap), len(heightMap[0])

        for i in range(m):
            heapq.heappush(heap, (heightMap[i][0], i, 0))
            heapq.heappush(heap, (heightMap[i][n - 1], i, n - 1))
            visited.add((i, 0))
            visited.add((i, n - 1))

        for j in range(n):
            heapq.heappush(heap, (heightMap[0][j], 0, j))
            heapq.heappush(heap, (heightMap[m - 1][j], m - 1, j))
            visited.add((0, j))
            visited.add((m - 1, j))

        res = 0
        while heap:
            height, i, j = heapq.heappop(heap)
            for x, y in [(i + 1, j), (i - 1, j), (i, j + 1), (i, j - 1)]:
                if 0 <= x < m and 0 <= y < n and (x, y) not in visited:
                    res += max(0, height - heightMap[x][y])
                    heapq.heappush(heap, (max(height, heightMap[x][y]), x, y))
                    visited.add((x, y))

        return res

复杂度分析:

  • 时间复杂度:O(mnlog(mn)),原因是我们需要遍历整个矩阵,然后每次插入都需要 log(mn) 的时间复杂度。
  • 空间复杂度:O(mn),原因是我们需要使用 heap 来存储数字。

295. Find Median from Data Stream

做一个平静且愿意倾听的人。
最后更新于 2023-05-25