用堆这种神奇的数据结构，在海量数据中瞬间找到你要的答案
堆（Heap）是一种特殊的完全二叉树，满足堆性质：每个节点的值都≥（或≤）其子节点的值。堆是优先队列的底层实现，也是TopK问题、中位数维护等场景的最优解法。Python中heapq模块默认实现小顶堆。
完全二叉树用数组存储，下标关系非常优雅：
这种表示不需要指针，内存连续，缓存友好。这是堆相比链式二叉树的最大优势之一。
import heapq
# 创建空堆并插入
heap = []
heapq.heappush(heap, 5) # O(log n) 插入
heapq.heappush(heap, 3)
heapq.heappush(heap, 1)
print(heap) # [1, 5, 3] — 注意不是完全排序
val = heapq.heappop(heap) # O(log n) 弹出最小值
print(val) # 1
# 从列表建堆 O(n)
nums = [5, 3, 1, 4, 2]
heapq.heapify(nums)
print(nums) # [1, 2, 3, 5, 4]
# 堆顶不弹出
print(nums[0]) # 1
# 大顶堆：取负数
max_heap = []
heapq.heappush(max_heap, -5)
heapq.heappush(max_heap, -3)
print(-heapq.heappop(max_heap)) # 5
# nsmallest / nlargest
print(heapq.nsmallest(2, [5,3,1,4,2])) # [1, 2]
print(heapq.nlargest(2, [5,3,1,4,2])) # [5, 4]| 场景 | 堆 | 排序 |
|---|---|---|
| 只需要TopK | O(n log k) ✅ | O(n log n) |
| 需要全部排序 | O(n log n) | O(n log n) ✅ |
| 动态插入+查询最值 | O(log n) 每次 ✅ | O(n log n) 每次重排 |
| 只需要最大/最小值 | O(1) 查询 ✅ | O(1) 查询但修改O(n) |
给定整数数组 nums 和整数 k，返回数组中第 k 个最大的元素。
class Solution:
def findKthLargest(self, nums: List[int], k: int) -> int:
# 小顶堆维护前k大的元素
heap = []
for num in nums:
heapq.heappush(heap, num)
if len(heap) > k:
heapq.heappop(heap) # 弹出最小的
return heap[0] # 堆顶即第k大
# 快速选择版本（进阶）
class Solution:
def findKthLargest(self, nums: List[int], k: int) -> int:
k = len(nums) - k # 第k大 = 第(n-k)小
def quick_select(l, r):
pivot = nums[r]
p = l
for i in range(l, r):
if nums[i] <= pivot:
nums[p], nums[i] = nums[i], nums[p]
p += 1
nums[p], nums[r] = nums[r], nums[p]
if p > k: return quick_select(l, p - 1)
elif p < k: return quick_select(p + 1, r)
else: return nums[p]
return quick_select(0, len(nums) - 1)| 方法 | 时间复杂度 | 空间复杂度 | 稳定性 |
|---|---|---|---|
| 小顶堆 | O(n log k) | O(k) | 稳定 ✅ |
| 快速选择 | O(n) 平均 | O(1) | 最坏O(n²) |
设计一个支持以下两种操作的数据结构：addNum(num) 添加一个数，findMedian() 返回当前所有数的中位数。
class MedianFinder:
def __init__(self):
self.small = [] # 大顶堆（存负数模拟）
self.large = [] # 小顶堆
def addNum(self, num: int) -> None:
# 先入大顶堆，再弹出最大入小顶堆
heapq.heappush(self.small, -num)
heapq.heappush(self.large, -heapq.heappop(self.small))
# 平衡：保证大顶堆元素 >= 小顶堆
if len(self.large) > len(self.small):
heapq.heappush(self.small, -heapq.heappop(self.large))
def findMedian(self) -> float:
if len(self.small) > len(self.large):
return -self.small[0]
return (-self.small[0] + self.large[0]) / 2
# 使用示例
mf = MedianFinder()
mf.addNum(1)
mf.addNum(2)
print(mf.findMedian()) # 1.5
mf.addNum(3)
print(mf.findMedian()) # 2.0addNum: O(log n) findMedian: O(1) 空间: O(n)
heapq.heappush(small, -num)，取出时再取负 -self.small[0]。这是面试常见陷阱，忘记取负会导致结果完全错误。给定整数数组 nums 和整数 k，返回出现频率前 k 高的元素。可以按任意顺序返回答案。
class Solution:
def topKFrequent(self, nums: List[int], k: int) -> List[int]:
# 第一步：统计频率
count = {}
for n in nums:
count[n] = count.get(n, 0) + 1
# 第二步：小顶堆维护TopK
heap = []
for num, freq in count.items():
heapq.heappush(heap, (freq, num))
if len(heap) > k:
heapq.heappop(heap)
return [num for freq, num in heap]
# 桶排序版本（O(n)时间）
class Solution:
def topKFrequent(self, nums: List[int], k: int) -> List[int]:
count = {}
for n in nums:
count[n] = count.get(n, 0) + 1
# 以频率为下标建桶
buckets = [[] for _ in range(len(nums) + 1)]
for num, freq in count.items():
buckets[freq].append(num)
# 从高到低收集
result = []
for i in range(len(buckets) - 1, -1, -1):
result.extend(buckets[i])
if len(result) >= k:
return result[:k]
return result| 方法 | 时间 | 空间 | 适用场景 |
|---|---|---|---|
| 堆 | O(n log k) | O(n) | 通用，k远小于n时最优 |
| 桶排序 | O(n) | O(n) | 频率范围有限时最优 |