Python——heapq模块-最大堆最小堆

由于queue不是Python标准库,所以在LeetCode等OJ上面不能直接使用,我们可以选择heapq来使用最大最小堆

使用示例

  • 堆排序示例

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    import heapq

    nums = [2, 3, 5, 1, 54, 23, 132]
    heap = []
    for num in nums:
        heapq.heappush(heap, num)
    # 等价于
    heapq.heapify(nums)

    # heap sort by incresing
    print([heapq.heappop(heap) for _ in range(len(nums))])

  • 加入元素

    1
    heapq.heappush(heap, num)

  • 弹出元素

    1
    num = heapq.heappop(heap)

  • 获取最大最小值

    1
    2
    3
    4
    5
    6
    7
    8
    9
    import heapq

    nums = [1, 3, 4, 5, 2]
    print(heapq.nlargest(3, nums))
    print(heapq.nsmallest(3, nums))

    #Output:
    [5, 4, 3]
    [1, 2, 3]

  • 获取堆顶元素

    1
    top = nums[0]

最大堆的实现

  • 由于Python heapq模块只实现了最小堆, 最大堆需要我们自己实现
  • 一种简单可行的实现方案:
    • 在加入和弹出时,把元素取反,从而实现最大堆