趋近智
堆和优先队列用于处理一个常见问题:从数据集中选出前 k 个元素,而无需对整个数据集进行排序。这是一项常遇到的任务,以多种形式出现,比如查找最常用的 k 个词、得分最高的 k 个预测,甚至作为更复杂算法的子程序。
假设你有一大串传感器读数,你只需要跟踪迄今为止最高的 5 个温度。每次新读数到达时都对所有读数进行排序会非常低效。堆提供了一种更好的方法。
主要思路是维护一个迄今为止遇到的 k 个最大元素的集合。最小堆出乎意料地适合这个任务。为什么是最小堆?因为它能让我们高效地找到并移除当前前 k 个元素中最小的那个。
算法如下:
heappush 将当前元素添加到堆中。heappop 移除最小元素(根部),并使用 heappush 添加当前元素。heapq 模块提供了一个方便的 heapreplace 函数,可以高效地合并这两个步骤。复杂度: 对于输入中的每个 n 元素,我们最多执行一次堆操作(推入或替换),这需要 O(logk) 时间,因为堆的大小受 k 限制。因此,总时间复杂度为 O(nlogk)。这比对整个数据集进行排序(O(nlogn))要好得多,尤其当 k 远小于 n 时。空间复杂度为 O(k),用于存储堆。
Python 实现 (heapq)
我们来使用 Python 的 heapq 模块查找列表中 3 个最大的元素。请记住,heapq 实现的是最小堆。
import heapq
def find_k_largest(data, k):
"""使用最小堆查找数据中 k 个最大的元素。"""
min_heap = [] # 我们的最小堆,用于存储 k 个最大元素
for element in data:
if len(min_heap) < k:
# 堆还没满,直接添加元素
heapq.heappush(min_heap, element)
else:
# 堆已满,将当前元素与堆中最小的元素(根部)进行比较
# 如果当前元素更大,则替换掉最小的元素
if element > min_heap[0]:
# 高效地弹出根部并推入新元素,同时保持堆的性质
heapq.heapreplace(min_heap, element)
# 堆现在包含 k 个最大元素
# min_heap[0] 是第 k 大的元素
# 返回排序后的列表,仅为清晰展示前 k 个元素
return sorted(min_heap, reverse=True)
# 示例用法
scores = [34, 56, 12, 89, 67, 99, 45, 78, 23, 91]
k = 3
top_k = find_k_largest(scores, k)
print(f"最大的 {k} 个分数是: {top_k}")
# 预期输出:最大的 3 个分数是: [99, 91, 89]
# 如果你只需要循环后的第 k 大元素:
# kth_largest = min_heap[0]
# print(f"第 {k} 大的分数是: {kth_largest}") # 应该是 89
查找 k 个最小元素遵循类似的思路,但这次我们使用最大堆。最大堆将存储迄今为止遇到的 k 个最小元素。当考虑一个新元素时,我们将其与堆中当前最大的元素(最大堆的根部)进行比较。如果新元素更小,它将替换堆中当前最大的元素。
由于 Python 的 heapq 只提供了最小堆的实现,我们使用一种常见技巧:在最小堆中存储元素的负值。这实际上将最小堆变成了关于原始值的最大堆。
-element)推入堆中。-element)与堆的根部(这是最小的负值,对应于 k 个最小元素中最大的原始值)进行比较。
-element 大于堆的根部(min_heap[0]),这意味着原始的 element 小于当前存储在我们 k 个最小元素集合中的最大值。使用 heapreplace 将根部替换为 -element。Python 实现(带负值处理的 heapq)
我们来使用这种技巧查找 4 个最小元素。
import heapq
def find_k_smallest(data, k):
"""使用模拟的最大堆查找数据中 k 个最小的元素。"""
# 使用存储负值的最小堆来模拟最大堆
max_heap = []
for element in data:
neg_element = -element # 对元素取负
if len(max_heap) < k:
heapq.heappush(max_heap, neg_element)
else:
# 将负值元素与根部(最小的负值 = 最大的原始值)进行比较
# 如果 neg_element > 根部,则 element < 堆中最大值
if neg_element > max_heap[0]:
heapq.heapreplace(max_heap, neg_element)
# 堆中包含 k 个最小元素的负值。取反回来。
# 取反回来并排序,以便清晰展示底部 k 个元素
smallest_k = sorted([-val for val in max_heap])
return smallest_k
# 示例用法
temps = [2.5, -1.0, 0.5, 3.0, -2.2, 1.8, 0.0, -0.5]
k = 4
bottom_k = find_k_smallest(temps, k)
print(f"最小的 {k} 个温度是: {bottom_k}")
# 预期输出:最小的 4 个温度是: [-2.2, -1.0, -0.5, 0.0]
# 循环后,第 k 小的元素是 -max_heap[0]
# 在此示例中,kth_smallest = -max_heap[0] 将是 0.0(因为 max_heap[0] 是 -0.0)
heapq.nlargest 和 heapq.nsmallestPython 的 heapq 模块为这些常见任务提供了直接函数:heapq.nlargest(k, iterable) 和 heapq.nsmallest(k, iterable)。这些函数内部实现了基于堆的选择算法,并提供了更简洁的接口。对于大多数实际应用,建议使用这些内置函数。
import heapq
data = [34, 56, 12, 89, 67, 99, 45, 78, 23, 91]
k = 3
largest_k = heapq.nlargest(k, data)
smallest_k = heapq.nsmallest(k, data)
print(f"Using nlargest: {largest_k}") # 输出:使用 nlargest: [99, 91, 89]
print(f"Using nsmallest: {smallest_k}") # 输出:使用 nsmallest: [12, 23, 34]
虽然自己实现逻辑有助于理解堆如何实现高效性,但在实际中,为简洁起见并可能获得更好的性能(因为有优化的 C 语言实现),应使用 nlargest 和 nsmallest。
在机器学习中,你可能会在哪些地方使用 k 选择技术?
理解如何使用堆高效地选择顶部或底部 k 个元素是一项有用的技能。它让你能够优化机器学习管道中涉及排序或选择的部分,避免在只需要极端值子集时进行完整排序的不必要开销。这项实践进一步说明了为当前任务选择正确数据结构所带来的效率提升。
这部分内容有帮助吗?
© 2026 ApX Machine Learning用心打造