Python：heapq模块使用

简介：heapq模块提供了堆队列算法的实现，也称为优先队列算法。

常用方法：
1、实现堆排序。
2、获取迭代器iter中n个最大的元素或者最小元素。

基础案例：

import heapq


def heap_sort(iterable):
    h = []
    for value in iterable:
        heapq.heappush(h, value)
    return [heapq.heappop(h) for i in range(len(h))]


if __name__ == '__main__':
    # 实现堆排序
    print(heap_sort([1, 3, 5, 7, 9, 2, 4, 6, 8, 0]))  # [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

    my_list = [2, 3, 7, 8, 1, 9, 7, 8, 5, 6]

    # 获取前 n 个最大值的列表
    n_large = heapq.nlargest(3, my_list)
    print("n_large", n_large)  # [9, 8, 8]

    # 获取前 n 个最小值
    n_small = heapq.nsmallest(3, my_list)
    print("n_small", n_small)  # [1, 2, 3]

进阶案例：

# 数组中的第 k大的数字
# 给定整数数组 nums 和整数 k，请返回数组中第 k 个最大的元素。
# 请注意，你需要找的是数组排序后的第 k 个最大的元素，而不是第 k 个不同的元素。
# 示例 1:
# 输入: [3,2,1,5,6,4] 和 k = 2
# 输出: 5
#
# 示例 2:
# 输入: [3,2,3,1,2,4,5,5,6] 和 k = 4
# 输出: 4


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


if __name__ == '__main__':
    s = Solution()
    nums = [3, 2, 3, 1, 2, 4, 5, 5, 6]
    r = s.findKthLargest(nums, 4)
    print(r)  # 4

微信公众号：玩转测试开发
欢迎关注，共同进步，谢谢！