最小堆最大堆基础概念与经典题目（Leetcode题解-Python语言）

cnhwl

已于 2022-04-25 16:18:01 修改

阅读量1k

点赞数 2

分类专栏： Python与Leetcode刷题文章标签： leetcode python 堆数据结构优先队列

于 2021-11-07 20:37:48 首次发布

本文链接：https://blog.csdn.net/cnhwl/article/details/121195814

版权

Python与Leetcode刷题专栏收录该内容

44 篇文章 5 订阅

订阅专栏

堆是一种特别的完全二叉树，符合以下两个定义即为堆：

1、完全二叉树；

2、每一个节点的值都必须大于等于或者小于等于其孩子节点的值。若是大于等于，即为最大堆，若是小于等于，即为最小堆。显然，最大堆的根节点是最大值，最小堆的根节点是最小值。

深度为 $k$ 的二叉树至多总共有 ${2^{k + 1}}$ 个节点（定义根节点所在深度 ${k_0}$ =0），节点数正好是 ${2^{k + 1}}$ 的二叉树就称为满二叉树；如果对满二叉树的节点编号，有二叉树的节点可以与编号一一对应的话，该二叉树就称为完全二叉树（又或者说，完全二叉树除最后一层外的其余层都是满的，并且最后一层要么是满的，要么在右边缺少连续若干节点）。

完全二叉树有两个特点：

1、具有 $n$ 个节点的完全二叉树的深度为 $\left\lfloor{{{\log }_2}n + 1}\right\rfloor$ （注： $\left\lfloor {} \right\rfloor$ 表示向下取整）

2、如果用数组表示完全二叉树，并且根节点存储在数组的索引1的位置的时候，任何一个节点的父节点索引位置为该节点的索引位置/2，任何一个节点的左孩子节点的索引位置为该节点的索引位置*2，任何一个节点的右孩子节点的索引位置为该节点的索引位置*2+1。

从零开始用Python实现最小堆的代码如下：

# 「最小堆」的实现
import sys

class MinHeap:
    def __init__(self, heapSize):
        # heapSize用于数组的大小，因为数组在创建的时候至少需要指明数组的元素个数
        self.heapSize = heapSize
        # 使用数组创建完全二叉树的结构，然后使用二叉树构建一个「堆」
        self.minheap = [0]*(heapSize+1)
        # realSize用于记录「堆」的元素个数
        self.realSize = 0

    #  添加元素函数
    def add(self, element):
        self.realSize += 1
        # 如果「堆」中元素的个数大于一开始设定的数组的个数，则返回「Add too many elements」
        if self.realSize > self.heapSize:
            print("Add too many elements!")
            self.realSize -= 1
            return
        # 将添加的元素添加到数组中
        self.minheap[self.realSize] = element
        # 新增元素的索引位置
        index = self.realSize
        # 新增元素的父节点的索引位置
        # 注意，如果用数组表示完全二叉树，并且根结点存储在数组的索引1的位置的时候，任何一个节点的父节点索引位置为「该节点的索引位置/2」，任何一个节点的左孩子节点的索引位置为「该节点的索引位置*2」，任何一个节点的右孩子节点的索引位置为「该节点的索引位置*2+1」
        parent = index // 2
        # 当添加的元素小于父节点时，需要将父节点的值和新增元素的值交换
        while (self.minheap[index] < self.minheap[parent] and index > 1):
            self.minheap[parent], self.minheap[index] = self.minheap[index], self.minheap[parent]
            index = parent
            parent = index // 2
    
    # 获取堆顶元素函数
    def peek(self):
        return self.minheap[1]
    
    # 删除堆顶元素函数
    def pop(self):
        # 如果当前「堆」的元素个数为0， 则返回「Don't have any element」
        if self.realSize < 1:
            print("Don't have any element!")
            return sys.maxsize
        else:
            # 当前「堆」中含有元素
            # self.realSize >= 1
            removeElement = self.minheap[1]
            # 将「堆」中的最后一个元素赋值给堆顶元素
            self.minheap[1] = self.minheap[self.realSize]
            self.realSize -= 1
            index = 1
            # 当删除的元素不是孩子节点时
            while (index < self.realSize and index <= self.realSize // 2):
                # 被删除节点的左孩子节点
                left = index * 2
                # 被删除节点的右孩子节点
                right = (index * 2) + 1
                # 当删除节点的元素大于 左孩子节点或者右孩子节点，代表该元素的值大，此时需要将该元素与左、右孩子节点中最小的值进行交换
                if (self.minheap[index] > self.minheap[left] or self.minheap[index] > self.minheap[right]):
                    if self.minheap[left] < self.minheap[right]:
                        self.minheap[left], self.minheap[index] = self.minheap[index], self.minheap[left]
                        index = left
                    else:
                        self.minheap[right], self.minheap[index] = self.minheap[index], self.minheap[right]
                        index = right
                else:
                    break
            return removeElement
    
    # 返回「堆」的元素个数
    def size(self):
        return self.realSize
    
    def toString(self):
        print(self.minheap[1 : self.realSize+1])
        

if __name__ == "__main__":
    	# 测试用例
        minHeap = MinHeap(5)
        minHeap.add(3)
        minHeap.add(1)
        minHeap.add(2)
        # [1,3,2]
        minHeap.toString()
        # 1
        print(minHeap.peek())
        # 1
        print(minHeap.pop())
        # 2
        print(minHeap.pop())
        # 3
        print(minHeap.pop())
        minHeap.add(4)
        minHeap.add(5)
        # [4,5]
        minHeap.toString()

实现最大堆的代码如下：

# 「最大堆」的实现
import sys

class MaxHeap:
    def __init__(self, heapSize):
        # heapSize用于数组的大小，因为数组在创建的时候至少需要指明数组的元素个数
        self.heapSize = heapSize
        # 使用数组创建完全二叉树的结构，然后使用二叉树构建一个「堆」
        self.maxheap = [0]*(heapSize+1)
        # realSize用于记录「堆」的元素个数
        self.realSize = 0

    #  添加元素函数
    def add(self, element):
        self.realSize += 1
        # 如果「堆」中元素的个数大于一开始设定的数组的个数，则返回「Add too many elements」
        if self.realSize > self.heapSize:
            print("Add too many elements!")
            self.realSize -= 1
            return
        # 将添加的元素添加到数组中
        self.maxheap[self.realSize] = element
        # 新增元素的索引位置
        index = self.realSize
        # 新增元素的父节点的索引位置
        # 注意，如果用数组表示完全二叉树，并且根结点存储在数组的索引1的位置的时候，任何一个节点的父节点索引位置为「该节点的索引位置/2」，任何一个节点的左孩子节点的索引位置为「该节点的索引位置*2」，任何一个节点的右孩子节点的索引位置为「该节点的索引位置*2+1」
        parent = index // 2
        # 当添加的元素大于父节点时，需要将父节点的值和新增元素的值交换
        while (self.maxheap[index] > self.maxheap[parent] and index > 1):
            self.maxheap[parent], self.maxheap[index] = self.maxheap[index], self.maxheap[parent]
            index = parent
            parent = index // 2
            
    # 获取堆顶元素函数
    def peek(self):
        return self.maxheap[1]
    
    # 删除堆顶元素函数
    def pop(self):
        # 如果当前「堆」的元素个数为0， 则返回「Don't have any element」
        if self.realSize < 1:
            print("Don't have any element!")
            return sys.maxsize
        else:
            # 当前「堆」中含有元素
            # self.realSize >= 1
            removeElement = self.maxheap[1]
            # 将「堆」中的最后一个元素赋值给堆顶元素
            self.maxheap[1] = self.maxheap[self.realSize]
            self.realSize -= 1
            index = 1
            # 当删除的元素不是孩子节点时
            while (index < self.realSize and index <= self.realSize // 2):
                # 被删除节点的左孩子节点
                left = index * 2
                # 被删除节点的右孩子节点
                right = (index * 2) + 1
                # 当删除节点的元素小于 左孩子节点或者右孩子节点，代表该元素的值小，此时需要将该元素与左、右孩子节点中最大的值进行交换
                if (self.maxheap[index] < self.maxheap[left] or self.maxheap[index] < self.maxheap[right]):
                    if self.maxheap[left] > self.maxheap[right]:
                        self.maxheap[left], self.maxheap[index] = self.maxheap[index], self.maxheap[left]
                        index = left
                    else:
                        self.maxheap[right], self.maxheap[index] = self.maxheap[index], self.maxheap[right]
                        index = right
                else:
                    break
            return removeElement
    
    # 返回「堆」的元素个数
    def size(self):
        return self.realSize
    
    def toString(self):
        print(self.maxheap[1 : self.realSize+1])
        

if __name__ == "__main__":
    	# 测试用例
        maxHeap = MaxHeap(5)
        maxHeap.add(1)
        maxHeap.add(2)
        maxHeap.add(3)
        # [3,1,2]
        maxHeap.toString()
        # 3
        print(maxHeap.peek())
        # 3
        print(maxHeap.pop())
        # 2
        print(maxHeap.pop())
        # 1
        print(maxHeap.pop())
        maxHeap.add(4)
        maxHeap.add(5)
        # [5,4]
        maxHeap.toString()

实现堆的关键是插入和删除，简单来说，以最小堆为例，插入操作就是把新的元素插入到二叉树的最后一个节点（保持完全二叉树），然后不断与其父节点比较大小，进行上移；删除操作就是把根节点元素与最后一个节点元素互换，删除最后一个节点（保持完全二叉树），然后根节点不断与其左右子节点比较大小，进行下移。

在Python中，已经内置了堆的实现，即标准库heapq，官方文档在此。由于官方只实现了最小堆，若想实现最大堆，只需要把元素取负即可。

最小堆示例：

# 最小堆完整代码
import heapq

# 新建一个列表
minHeap = []
# 将列表堆化，即将列表转换为最小堆
heapq.heapify(minHeap)
# 分别往最小堆中添加3，1，2
heapq.heappush(minHeap, 3)
heapq.heappush(minHeap, 1)
heapq.heappush(minHeap, 2)
# 查看最小堆的所有元素，结果为：[1,3,2]
print("minHeap: ",minHeap)
# 获取最小堆的堆顶元素
peekNum = minHeap[0]
# 结果为：1
print("peek number: ", peekNum)
# 删除最小堆的堆顶元素
popNum = heapq.heappop(minHeap)
# 结果为：1
print("pop number: ", popNum)
# 查看删除1后最小堆的堆顶元素，结果为：2
print("peek number: ", minHeap[0])
# 查看最小堆的所有元素，结果为：[2,3]
print("minHeap: ",minHeap)
# 获取堆的元素个数，即堆的长度
size = len(minHeap)
# 结果为：2
print("minHeap size: ", size)

最大堆示例：

# 最大堆完整代码
import heapq

# 新建一个列表
maxHeap = []
# 将列表堆化，此时的堆是最小堆，我们需要将元素取反技巧，将最小堆转换为最大堆
heapq.heapify(maxHeap)
# 分别往堆中添加1，3，2，注意此时添加的是-1，-3，-2，原因是需要将元素取反，最后将最小堆转换为最大堆
heapq.heappush(maxHeap, 1*-1)
heapq.heappush(maxHeap, 3*-1)
heapq.heappush(maxHeap, 2*-1)
# 查看堆中所有元素：[-3, -1, -2]
print("maxHeap: ",maxHeap)
# 查看堆中的最大元素，即当前堆中最小值*-1
peekNum = maxHeap[0]
# 结果为：3
print("peek number: ", peekNum*-1)
# 删除堆中最大元素，即当前堆中最小值
popNum = heapq.heappop(maxHeap)
# 结果为：3
print("pop number: ", popNum*-1)
# 查看删除3后堆中最大值， 结果为：2
print("peek number: ", maxHeap[0]*-1)
# 查看堆中所有元素，结果为：[-2,-1]
print("maxHeap: ",maxHeap)
# 查看堆的元素个数，即堆的大小
size = len(maxHeap)
# 结果为：2
print("maxHeap size: ", size)

简单练手题

1046. 最后一块石头的重量

class Solution:
    def lastStoneWeight(self, stones: List[int]) -> int:
        heap = [-i for i in stones]
        heapq.heapify(heap)
        while heap:
            stone1 = heapq.heappop(heap) * -1
            if not heap:
                return stone1
            else:
                stone2 = heapq.heappop(heap) * -1
                stone1 = stone1 - stone2
                heapq.heappush(heap, stone1 * -1)
        return 0

只需要创建一个最大堆，每次弹出两个石头，进行相减后放回堆中，若只剩一个石头则返回该石头，否则返回0。

经典题目

最经典的一类题目莫过于 Top K 和 The Kth ，即求数组（大小为 N）中最大或最小的 K 个数或者第 K 个数。这类问题一般有两种思路（以求取最小的 K 个数或者第 K 个数为例）：

1、创建一个大小为 N 的最小堆，然后对其进行 K 次弹出（heappop）操作，由于每次都是弹出最小值，所以得到的结果一定就是最小的 K 个数，只要最小的第 K 个数也可以。时间复杂度是 $O(K\log N)$ ，是因为进行了 K 次弹出操作，而每次弹出后最小堆都会比较 $\log N$ 次把下一个最小值放到根节点，因此是 $O(K\log N)$ ，空间复杂度则为 $O (N)$ 。

2、创建一个大小为 K 的最大堆，遍历数组（N 次遍历），首先数组顺序的前 K 个元素加入最大堆（填满），然后当最大堆的元素个数达到 K 时，后面的遍历就要将当前遍历元素与堆顶元素进行比较，如果当前元素大于堆顶元素，则放弃当前元素，继续遍历下一个元素；如果当前元素小于堆顶元素，则删除堆顶元素，将当前元素加入到最大堆中（heapreplace）。最后得到的最大堆中的 K 个元素就是最小的 K 个元素。时间复杂度是 $O(N\log K)$ ，是因为进行了 N 次遍历，而每次遍历元素若加入最大堆，就会比较 $\log K$ 次把下一个最小值放到根节点，因此是 $O(N\log K)$ ，空间复杂度则为 $O (K)$ 。

剑指 Offer 40. 最小的k个数（面试题 17.14. 最小K个数）

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        heapq.heapify(arr)
        ans = []
        for _ in range(k):
            ans.append(heapq.heappop(arr))
        return ans

建立最小堆，弹出K个数即为最小的K个数。

215. 数组中的第K个最大元素（剑指 Offer II 076. 数组中的第 k 大的数字）

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        heapq.heapify(nums)
        while len(nums) > k:
            heapq.heappop(nums)
        return nums[0]

此题用取负实现最大堆也可以，但是更好的是实现一个长度为K的最小堆，则此时堆的最小值就是数组中的第K个最大元素。

703. 数据流中的第 K 大元素（剑指 Offer II 059. 数据流的第 K 大数值）

class KthLargest:

    def __init__(self, k: int, nums: List[int]):
        self.k = k
        self.heap = nums
        heapq.heapify(self.heap)

    def add(self, val: int) -> int:
        heapq.heappush(self.heap, val)
        while len(self.heap) > self.k:
            heapq.heappop(self.heap)
        return self.heap[0]

思路与上一题一样，实现一个长度为K的最小堆，即可得到数组中第K大元素。

1985. 找出数组中的第 K 大整数

class Solution:
    def kthLargestNumber(self, nums: List[str], k: int) -> str:
        heap = [int(i) for i in nums]
        heapq.heapify(heap)
        while len(heap) > k:
            heapq.heappop(heap)
        return str(heap[0])

还是第K大的数，加上了字符串到整数的转换。

347. 前 K 个高频元素（剑指 Offer II 060. 出现频率最高的 k 个数字）

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

使用 collections.Counter()统计频率，然后建立长度为K的最小堆，若频率值比堆的最小值大，则替换掉它。此处有两个知识点要注意：

1、堆元素可以为元组，但比较值必须在前面，即 (val, key)。
2、heapq.heappushpop(heap, item)是先将 item 放入堆中，然后弹出并返回 heap 的最小元素；
heapq.heapreplace(heap, item)是先弹出并返回 heap 中最小的一项，然后推入新的 item，相当于poppush。

692. 前K个高频单词

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

此题若沿用上一题思路不太好做，这里改为建立最大堆，直接弹出K个元素即可，隐含的机制是堆会自动对元组 (-value, key) 进行从左到右优先级的排序。

451. 根据字符出现频率排序

class Solution:
    def frequencySort(self, s: str) -> str:
        count = collections.Counter(list(s))
        heap = []
        for key, value in count.items():
            heapq.heappush(heap, (-value, key))
        ans = ''
        for _ in range(len(heap)):
            value, key = heapq.heappop(heap)
            ans = ans + key * -value
        return ans

这题也是统计词频并排序，只是输出需要改动一下而已。

373. 查找和最小的K对数字（剑指 Offer II 061. 和最小的 k 个数对）

class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        heap = []
        for n1 in nums1:
            for n2 in nums2:
                if len(heap) < k:
                    heapq.heappush(heap, (-n1-n2, [n1, n2]))
                elif heap and -heap[0][0] > n1 + n2:
                    heapq.heapreplace(heap, (-n1-n2, [n1, n2]))
                else:
                    break
        return [i[1] for i in heap]

本题不能用最小堆，否则每次比较只能与最小值比较，不能替换掉其他值。所以用长度为K的最大堆，与347题基本一样。

263. 丑数

class Solution:
    def isUgly(self, n: int) -> bool:
        if n <= 0:
            return False
        factors = [2, 3, 5]
        for factor in factors:
            while n % factor == 0:
                n //= factor
        return n == 1

数学类型的题目，基本写法记住就好。

264. 丑数 II（剑指 Offer 49. 丑数）（面试题 17.09. 第 k 个数）

class Solution:
    def nthUglyNumber(self, n: int) -> int:
        factors = [2, 3, 5]
        seen = {1}
        heap = [1]
        for i in range(n-1):
            cur = heapq.heappop(heap)
            for factor in factors:
                nxt = cur * factor
                if nxt not in seen:
                    seen.add(nxt)
                    heapq.heappush(heap, nxt)
        return heapq.heappop(heap)

此题不是判断丑数，而是寻找第 N 个丑数，那就使用一个最小堆，每次弹出最小的丑数，然后将其与3个因子结合可以生成3个新的丑数，并加入到最小堆中，重复 N-1 次，最后返回第 N 个即为所求。

378. 有序矩阵中第 K 小的元素

class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        n = len(matrix)
        heap = [(matrix[i][0], i, 0) for i in range(n)]
        heapq.heapify(heap)

        for _ in range(k-1):
            num, x, y = heapq.heappop(heap)   # x是在哪一行，y是一行中的哪个位置（列）
            if y != n - 1:
                heapq.heappush(heap, (matrix[x][y + 1], x, y + 1))
        
        return heapq.heappop(heap)[0]

本题是堆与矩阵的结合题，用二分查找是最优解法，但此处还是使用了堆。只需要用一个最小堆记录最小值，每弹出一个最小值就把它在矩阵中右边的元素加入到堆中，若右边没有元素则跳过，重复 K 次即可。

1439. 有序矩阵中的第 k 个最小数组和

class Solution:
    def kthSmallest(self, mat, k: int) -> int:
        m = len(mat)
        n = len(mat[0])
        heap = []
        cur_sum = 0
        # 第一列的和
        for i in range(m):
            cur_sum += mat[i][0]
        # 各行的指针
        pointers = [0] * m
        heapq.heappush(heap, [cur_sum, tuple(pointers)])
        # 出现过的指针组合放入seen
        seen = set()
        seen.add(tuple(pointers))  # 必须用tuple才能hash，才能放入集合

        for _ in range(k-1):
            # 从堆中pop出cur_sum(最小数组和)和pointers(指针数组)
            cur_sum, pointers = heapq.heappop(heap)
            # 每个指针轮流后移一位，将new_sum(新的数组和)和new_pointers(新的指针数组)push入堆
            for idx, pointer in enumerate(pointers):
                if pointer < n - 1:
                    # tuple变为list修改再变回tuple
                    new_pointers = list(pointers)
                    new_pointers[idx] = pointer + 1
                    new_pointers = tuple(new_pointers)
                    if new_pointers not in seen:
                        new_sum = cur_sum - mat[idx][pointer] + mat[idx][pointer + 1]
                        heapq.heappush(heap, [new_sum, new_pointers])
                        seen.add(new_pointers)
                        
        return heapq.heappop(heap)[0]

这道困难题可以借鉴丑数的思路，用最小堆记录和值 sum 与指针组合 pointers（注意pointers必须用元组，否则不能哈希，放不进集合），然后用集合 seen 记录出现过的 pointers，如果没有出现过，则 push 进最小堆并记录到 seen 中。重复 K 次的 pop，然后让每行的指针都后移一位，直到无法移动为止。

cnhwl

关注

2
点赞
踩
4

收藏

觉得还不错? 一键收藏
0
评论
最小堆最大堆基础概念与经典题目（Leetcode题解-Python语言）

堆是一种特别的完全二叉树，符合以下两个定义即为堆：1、完全二叉树；2、每一个节点的值都必须大于等于或者小于等于其孩子节点的值。若是大于等于，即为最大堆，若是小于等于，即为最小堆。显然，最大堆的根节点是最大值，最小堆的根节点是最小值。深度为kkk的二叉树至多总共有2k+1{2^{k + 1}}2k+1个节点（定义根节点所在深度k0{k_0}k0=0），节点数正好是2k+1{2^{k + 1}}2k+1的二叉树就称为满二叉树；如果对满二叉树的节点编号，有二叉树的节点可以与编号一一对应的话，该二叉树就称为
复制链接

扫一扫