python

关注公众号 jb51net

关闭
首页 > 脚本专栏 > python > Python heapq堆操作

Python heapq堆操作全解析

作者：老师好，我是刘同学

本文主要介绍了Python heapq堆操作全解析,文中通过示例代码介绍的非常详细,对大家的学习或者工作具有一定的参考学习价值,需要的朋友们下面随着小编来一起学习学习吧

1. heapq 库概述

Python 的 heapq 库是基于堆数据结构实现的标准库模块，它提供了对小顶堆（min-heap）的高效操作支持。堆是一种特殊的完全二叉树结构，其中父节点的值总是小于或等于其所有子节点的值（小顶堆特性）。该库的时间复杂度为 O(log n)，在需要频繁插入和删除最小元素的场景下表现出色。

2. 核心函数详解

2.1 基础堆操作函数

函数名功能描述时间复杂度使用场景
heapify(x)将列表 x 转换为堆结构O(n)列表初始化堆
heappush(heap, item)向堆中插入新元素O(log n)动态添加元素
heappop(heap)弹出并返回最小元素O(log n)获取最小元素
heapreplace(heap, item)弹出最小元素并插入新元素O(log n)替换堆顶元素
heappushpop(heap, item)先插入再弹出最小元素O(log n)高效插入弹出

代码示例：基础堆操作

import heapq
# 初始化列表
data = [3, 1, 4, 1, 5, 9, 2, 6]
# 将列表转换为堆（原地操作）
heapq.heapify(data)
print(f"堆化后的列表: {data}")  # 输出: [1, 1, 2, 3, 5, 9, 4, 6]
# 向堆中插入元素
heapq.heappush(data, 0)
print(f"插入0后的堆: {data}")  # 输出: [0, 1, 2, 1, 5, 9, 4, 6, 3]
# 弹出最小元素
min_element = heapq.heappop(data)
print(f"弹出的最小元素: {min_element}")  # 输出: 0
print(f"弹出后的堆: {data}")  # 输出: [1, 1, 2, 3, 5, 9, 4, 6]

2.2 批量查询函数

函数名功能描述时间复杂度适用场景
nlargest(n, iterable)返回前n个最大元素O(n log k)Top-K 最大元素
nsmallest(n, iterable)返回前n个最小元素O(n log k)Top-K 最小元素

代码示例：Top-K 问题解决

import heapq
import random

# 生成测试数据
numbers = [random.randint(1, 1000) for _ in range(100)]

# 获取最大的5个元素
largest_5 = heapq.nlargest(5, numbers)
print(f"最大的5个元素: {largest_5}")

# 获取最小的5个元素
smallest_5 = heapq.nsmallest(5, numbers)
print(f"最小的5个元素: {smallest_5}")

# 使用key参数进行自定义比较
words = ['apple', 'banana', 'cherry', 'date', 'elderberry']
longest_3 = heapq.nlargest(3, words, key=len)
print(f"最长的3个单词: {longest_3}")  # 输出: ['elderberry', 'banana', 'cherry']

2.3 高级操作函数

heapq.merge(*iterables) 函数用于合并多个已排序的输入序列，返回一个排序后的迭代器。

import heapq

# 合并多个有序序列
list1 = [1, 3, 5, 7]
list2 = [2, 4, 6, 8]
list3 = [0, 9, 10]

merged = list(heapq.merge(list1, list2, list3))
print(f"合并后的有序列表: {merged}")  # 输出: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

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):
        """弹出优先级最高的元素"""
        if self._heap:
            return heapq.heappop(self._heap)[-1]
        raise IndexError("优先级队列为空")
    
    def is_empty(self):
        return len(self._heap) == 0

# 使用示例
pq = PriorityQueue()
pq.push("任务A", 3)
pq.push("任务B", 1)  # 最高优先级
pq.push("任务C", 2)

while not pq.is_empty():
    print(f"执行: {pq.pop()}")
# 输出: 执行: 任务B → 执行: 任务C → 执行: 任务A

3.2 实时数据流的中位数查找

import heapq

class MedianFinder:
    def __init__(self):
        # 最大堆（使用负数模拟）和最小堆
        self.max_heap = []  # 存储较小的一半
        self.min_heap = []  # 存储较大的一半
    
    def add_num(self, num):
        if not self.max_heap or num <= -self.max_heap[0]:
            heapq.heappush(self.max_heap, -num)
        else:
            heapq.heappush(self.min_heap, num)
        
        # 平衡两个堆
        if len(self.max_heap) > len(self.min_heap) + 1:
            heapq.heappush(self.min_heap, -heapq.heappop(self.max_heap))
        elif len(self.min_heap) > len(self.max_heap):
            heapq.heappush(self.max_heap, -heapq.heappop(self.min_heap))
    
    def find_median(self):
        if len(self.max_heap) == len(self.min_heap):
            return (-self.max_heap[0] + self.min_heap[0]) / 2
        else:
            return -self.max_heap[0]

# 使用示例
finder = MedianFinder()
for num in [1, 3, 2, 6, 4, 5]:
    finder.add_num(num)
    print(f"当前中位数: {finder.find_median()}")

3.3 堆排序算法

import heapq

def heap_sort(iterable):
    """使用堆排序算法对可迭代对象进行排序"""
    heap = list(iterable)
    heapq.heapify(heap)  # 构建最小堆
    return [heapq.heappop(heap) for _ in range(len(heap))]

# 排序示例
unsorted_data = [9, 2, 7, 5, 1, 8, 3, 6, 4]
sorted_data = heap_sort(unsorted_data)
print(f"堆排序结果: {sorted_data}")  # 输出: [1, 2, 3, 4, 5, 6, 7, 8, 9]

4. 高级技巧与性能优化

4.1 实现最大堆

由于 heapq 默认实现的是最小堆，可以通过存储负值来模拟最大堆：

import heapq

class MaxHeap:
    def __init__(self):
        self._heap = []
    
    def push(self, item):
        heapq.heappush(self._heap, -item)
    
    def pop(self):
        return -heapq.heappop(self._heap)
    
    def peek(self):
        return -self._heap[0] if self._heap else None

# 最大堆使用示例
max_heap = MaxHeap()
for num in [3, 1, 4, 1, 5]:
    max_heap.push(num)

print("最大堆元素弹出顺序:")
while max_heap._heap:
    print(max_heap.pop())
# 输出: 5, 4, 3, 1, 1

4.2 自定义对象堆操作

import heapq

class Task:
    def __init__(self, name, priority, duration):
        self.name = name
        self.priority = priority
        self.duration = duration
    
    def __lt__(self, other):
        # 定义比较规则：优先级高的在前，相同优先级时持续时间短的在前
        if self.priority == other.priority:
            return self.duration < other.duration
        return self.priority > other.priority
    
    def __repr__(self):
        return f"Task({self.name}, priority:{self.priority}, duration:{self.duration})"

# 自定义对象堆操作
tasks = [
    Task("紧急任务", 3, 2),
    Task("普通任务", 1, 5),
    Task("重要任务", 2, 3)
]

heap = []
for task in tasks:
    heapq.heappush(heap, task)

print("任务执行顺序:")
while heap:
    print(heapq.heappop(heap))

5. 性能对比与最佳实践

5.1 不同场景下的性能选择

操作场景推荐方法时间复杂度优势
一次性获取Top-Knlargest()/nsmallest()O(n log k)代码简洁
持续插入和弹出heappush() + heappop()O(log n)动态高效
多个有序序列合并heapq.merge()O(n log k)内存友好

5.2 内存优化技巧

import heapq

# 流式处理大数据集
def process_large_dataset(data_stream, top_n=10):
    """使用堆处理大数据流，只维护Top-N元素"""
    heap = []
    
    for item in data_stream:
        if len(heap) < top_n:
            heapq.heappush(heap, item)
        elif item > heap[0]:  # 对于最大Top-N，使用最小堆
            heapq.heapreplace(heap, item)
    
    return sorted(heap, reverse=True)

# 模拟大数据流处理
import random
data_stream = (random.randint(1, 10000) for _ in range(100000))
top_10 = process_large_dataset(data_stream, 10)
print(f"大数据流中的Top-10: {top_10}")

Python 的 heapq 库通过提供高效的堆操作函数，在算法优化、数据处理和系统设计等多个领域发挥着重要作用。掌握这些函数的正确使用方法和适用场景，能够显著提升程序的性能和代码的可维护性。

到此这篇关于Python heapq堆操作全解析的文章就介绍到这了,更多相关Python heapq堆操作内容请搜索脚本之家以前的文章或继续浏览下面的相关文章希望大家以后多多支持脚本之家！

您可能感兴趣的文章:
阅读全文