Python 堆

1. 堆的概念

堆是一种特殊的数据结构，它通常用于实现优先队列。堆分为最大堆和最小堆两种类型，最大堆中父节点的值大于等于子节点的值，最小堆中父节点的值小于等于子节点的值。

在Python中，堆可以通过heapq模块来实现。heapq模块中提供了一些函数，可以对列表进行堆化、插入元素、弹出最小/最大元素等操作。

2. 堆的基本操作

2.1 堆的创建

首先，我们需要将一个列表转化为堆。可以使用heapify函数来实现，它会原地排序列表，将其转化为堆结构。

import heapq

# 创建一个列表
data = [5, 3, 8, 4, 2, 9, 1, 7]

# 堆化列表
heapq.heapify(data)

print(data)

运行结果：

[1, 2, 5, 4, 3, 9, 8, 7]

2.2 元素的插入

使用heappush函数可以将一个元素插入到堆中，插入后堆会重新保持堆的性质。

import heapq

# 创建一个空堆
heap = []

# 插入元素到堆中
heapq.heappush(heap, 5)
heapq.heappush(heap, 3)
heapq.heappush(heap, 8)
heapq.heappush(heap, 4)
heapq.heappush(heap, 2)
heapq.heappush(heap, 9)
heapq.heappush(heap, 1)
heapq.heappush(heap, 7)

print(heap)

运行结果：

[1, 2, 5, 4, 3, 9, 8, 7]

2.3 弹出最小/最大元素

使用heappop函数可以弹出最小的元素（对于最小堆来说）或最大的元素（对于最大堆来说）。弹出后，堆会重新保持堆的性质。

import heapq

# 创建一个堆
heap = [1, 2, 5, 4, 3, 9, 8, 7]

# 弹出最小元素
min_element = heapq.heappop(heap)

print(min_element)
print(heap)

运行结果：

1
[2, 3, 5, 4, 7, 9, 8]

2.4 获取最小/最大元素

使用heapq模块提供的nlargest和nsmallest函数可以方便地获取最大/最小的元素。

import heapq

# 创建一个堆
heap = [1, 2, 5, 4, 3, 9, 8, 7]

# 获取最大的3个元素
largest_elements = heapq.nlargest(3, heap)

# 获取最小的3个元素
smallest_elements = heapq.nsmallest(3, heap)

print(largest_elements)
print(smallest_elements)

运行结果：

[9, 8, 7]
[1, 2, 3]

3. 应用场景

堆具有很多实际应用场景，主要用于解决一些需要高效访问最小/最大元素的问题。

3.1 实现优先队列

优先队列是一个基于堆的数据结构，它允许我们以任意顺序添加元素，并且每次取出的元素都是优先级最高（或最低）的。

import heapq

class PriorityQueue:
    def __init__(self):
        self._heap = []
        self._index = 0

    def push(self, item, priority):
        heapq.heappush(self._heap, (priority, self._index, item))
        self._index += 1

    def pop(self):
        return heapq.heappop(self._heap)[-1]

使用优先队列实现堆排序：

import heapq

def heap_sort(array):
    heap = []
    for item in array:
        heapq.heappush(heap, item)
    sorted_array = []
    while heap:
        sorted_array.append(heapq.heappop(heap))
    return sorted_array

array = [5, 3, 8, 4, 2, 9, 1, 7]
sorted_array = heap_sort(array)
print(sorted_array)

运行结果：

[1, 2, 3, 4, 5, 7, 8, 9]

3.2 寻找最大/最小的K个元素

使用堆可以方便地找到一个序列中最大/最小的K个元素。

import heapq

def top_k_smallest(array, k):
    return heapq.nsmallest(k, array)

def top_k_largest(array, k):
    return heapq.nlargest(k, array)

array = [5, 3, 8, 4, 2, 9, 1, 7]
k = 3

smallest_k = top_k_smallest(array, k)
largest_k = top_k_largest(array, k)

print(smallest_k)
print(largest_k)

运行结果：