python堆排序之heapq

python 的 heapq模块提供了堆排序实现，具体如下：

创建堆：

• 只有最小堆，没有最大堆
• 想要有最大堆的效果，可以将数据取相反数
• heapq.heappush
• heapq.heapify

访问堆内容

• heapq.heappop
• heapq.nlargest （最小堆）
• heapq.nsmallest （最大堆）

更新/替换堆

• heapq.heaprepalce

1、heapq.heappush & heapq.heappop

import heapq
nums = [21, 13, 52, 16, 34, 21, 119]
heap = []
for num in nums:
    heapq.heappush(heap, num)
    
print(heap[0])  # 弹出最小值，使用heap[0]
print([heapq.heappop(heap) for _ in range(len(nums))])  # 堆排序结果

结果：
13
[13, 16, 21, 21, 34, 52, 119]

2、heapq.heapify

import heapq
nums = [21, 13, 52, 16, 34, 21, 119]
heapq.heapify(nums)
    
print(nums[0])  # 弹出最小值，使用nums[0]
print([heapq.heappop(nums) for _ in range(len(nums))])  # 堆排序结果

结果：
13
[13, 16, 21, 21, 34, 52, 119]

3、 heapq.heappop & heapq.nlargest & heapq.nsmallest

import heapq
nums = [21, 13, 52, 16, 34, 21, 119]

print(heapq.nlargest(3, nums))
print(heapq.nsmallest(3, nums))
结果：
[119, 52, 34]
[13, 16, 21]

这两个函数还接受一个参数key，用于处理dict等类型数据

score= [
    {'name': 'xiaoming', 'score': 100},
    {'name': 'xiaozhang', 'score': 90},
    {'name': 'xiaoli', 'score': 92},
    {'name': 'xiaowang', 'score': 65},
    {'name': 'xiaosun', 'score': 78}
]
print(heapq.nsmallest(3, score, key=lambda s: s['score']))
print(heapq.nlargest(3, score, key=lambda s: s['score']))

结果：
[{‘score’: 65, ‘name’: ‘xiaowang’}, {‘score’: 78, ‘name’: ‘xiaosun’}, {‘score’: 90, ‘name’: ‘xiaozhang’}]
[{‘score’: 100, ‘name’: ‘xiaoming’}, {‘score’: 92, ‘name’: ‘xiaoli’}, {‘score’: 90, ‘name’: ‘xiaozhang’}]

4、 heapq.heaprepalce

import heapq

nums = [21, 13, 52, 16, 34, 21, 119]
heapq.heapify(nums)

heapq.heapreplace(nums, 23)
print([heapq.heappop(nums) for _ in range(len(nums))])

结果：
[16, 21, 21, 23, 34, 52, 119]

大顶堆pyhton实现

class Solution:
    # 调整为大顶堆
    def heapify(self, arr, index, end):
        left = index * 2 + 1
        right = left + 1
        while left <= end:
            # 当前节点为非叶子节点
            max_index = index
            if arr[left] > arr[max_index]:
                max_index = left
            if right <= end and arr[right] > arr[max_index]:
                max_index = right
            if index == max_index:
                # 如果不用交换，则说明已经交换结束
                break
            arr[index], arr[max_index] = arr[max_index], arr[index]
            # 继续调整子树
            index = max_index
            left = index * 2 + 1
            right = left + 1

    # 初始化大顶堆
    def buildMaxHeap(self, arr):
        size = len(arr)
        # (size-2) // 2 是最后一个非叶节点，叶节点不用调整
        for i in range((size - 2) // 2, -1, -1):
            self.heapify(arr, i, size - 1)
        return arr

    # 升序堆排序，思路如下：
    # 1. 先建立大顶堆
    # 2. 让堆顶最大元素与最后一个交换，然后调整第一个元素到倒数第二个元素，这一步获取最大值
    # 3. 再交换堆顶元素与倒数第二个元素，然后调整第一个元素到倒数第三个元素，这一步获取第二大值
    # 4. 以此类推，直到最后一个元素交换之后完毕。
    def maxHeapSort(self, arr):
        self.buildMaxHeap(arr)
        size = len(arr)
        for i in range(size):
            arr[0], arr[size-i-1] = arr[size-i-1], arr[0]
            self.heapify(arr, 0, size-i-2)
        return arr

    def sortArray(self, nums: List[int]) -> List[int]:
        return self.maxHeapSort(nums)

Ref：

  [1].leetcode
  [2].python 模块：heapq—堆及堆排序