堆

本章目标：理解堆的原理，手写堆结构，掌握 Top-K 问题模板。

你将学到

• 完全二叉树与数组表示的关系
• 大根堆与小根堆
• 上浮（sift up）与下沉（sift down）
• 堆排序
• Top-K 问题模板
• Swift 中手写 Priority Queue

1. 概念

堆是一种特殊的完全二叉树：大根堆中每个节点的值 >= 其子节点的值，小根堆则 <=。堆通常用数组实现，因为完全二叉树的节点可以映射到连续的数组下标：

• 节点 i 的左子节点：2*i + 1
• 节点 i 的右子节点：2*i + 2
• 节点 i 的父节点：(i - 1) / 2

堆的核心应用：优先队列、Top-K 问题、堆排序、合并 K 个有序序列。Swift 标准库没有内置堆，需要自己实现或使用 Swift Collections 的 Heap。

2. 核心模板

小根堆实现

struct MinHeap {
    private var heap: [Int] = []

    var isEmpty: Bool { heap.isEmpty }
    var count: Int { heap.count }
    var peek: Int? { heap.first }

    func parentIndex(_ i: Int) -> Int { (i - 1) / 2 }
    func leftChildIndex(_ i: Int) -> Int { 2 * i + 1 }
    func rightChildIndex(_ i: Int) -> Int { 2 * i + 2 }

    mutating func insert(_ val: Int) {
        heap.append(val)
        siftUp(heap.count - 1)
    }

    private mutating func siftUp(_ index: Int) {
        var i = index
        while i > 0 && heap[i] < heap[parentIndex(i)] {
            heap.swapAt(i, parentIndex(i))
            i = parentIndex(i)
        }
    }

    mutating func extractMin() -> Int? {
        guard !heap.isEmpty else { return nil }
        let minVal = heap[0]
        heap[0] = heap[heap.count - 1]
        heap.removeLast()
        if !heap.isEmpty {
            siftDown(0)
        }
        return minVal
    }

    private mutating func siftDown(_ index: Int) {
        var i = index
        while true {
            var smallest = i
            let l = leftChildIndex(i), r = rightChildIndex(i)
            if l < heap.count && heap[l] < heap[smallest] { smallest = l }
            if r < heap.count && heap[r] < heap[smallest] { smallest = r }
            if smallest == i { break }
            heap.swapAt(i, smallest)
            i = smallest
        }
    }
}

复杂度

• 插入/删除：O(log n)
• 查看堆顶：O(1)
• 建堆：O(n)

3. 经典题目精讲

LeetCode 215 · 数组中第 K 个最大元素

题意：找数组中第 K 大的元素。

思路：维护大小为 K 的小根堆，遍历完数组后堆顶就是第 K 大。也可用快速选择算法 O(n) 平均。

func findKthLargest(_ nums: [Int], _ k: Int) -> Int {
    var heap = MinHeap()
    for num in nums {
        heap.insert(num)
        if heap.count > k {
            _ = heap.extractMin()
        }
    }
    return heap.peek!
}

LeetCode 347 · 前 K 个高频元素

题意：返回数组中出现频率前 K 高的元素。

思路：先用哈希表统计频率，再用小根堆按频率保留前 K 个。

func topKFrequent(_ nums: [Int], _ k: Int) -> [Int] {
    // 1. 统计频率
    var freq: [Int: Int] = [:]
    for num in nums { freq[num, default: 0] += 1 }

    // 2. 用数组模拟最小堆（按频率排序）
    var heap: [(num: Int, count: Int)] = []
    for (num, count) in freq {
        heap.append((num, count))
        heap.sort { $0.count < $1.count }  // 简化写法，正式用堆
        if heap.count > k { heap.removeFirst() }
    }
    return heap.map { $0.num }
}

LeetCode 295 · 数据流中位数

题意：动态添加数字，随时能 O(1) 获取当前中位数。

思路：大根堆存较小的一半，小根堆存较大的一半，保持两堆大小差不超过 1。

class MedianFinder {
    var leftHeap: [Int] = []   // 大根堆（存负数模拟）
    var rightHeap: [Int] = []  // 小根堆

    func addNum(_ num: Int) {
        // 先加入大根堆
        leftHeap.append(-num)
        leftHeap.sort()
        leftHeap.reverse()
        // 把大根堆最大值移到小根堆
        let val = -leftHeap.removeFirst()
        rightHeap.append(val)
        rightHeap.sort()
        // 平衡：小根堆不能比大根堆多
        if rightHeap.count > leftHeap.count {
            let v = rightHeap.removeFirst()
            leftHeap.append(-v)
            leftHeap.sort()
            leftHeap.reverse()
        }
    }

    func findMedian() -> Double {
        if leftHeap.count > rightHeap.count {
            return Double(-leftHeap[0])
        }
        return Double(-leftHeap[0] + rightHeap[0]) / 2.0
    }
}

4. 易错点

• ⚠️ Swift 没有内置堆，正式刷题建议引入 Swift Collections 或手写
• ⚠️ 数组下标从 0 开始时，父子关系与从 1 开始不同
• ⚠️ 堆排序是原地排序，空间 O(1)，但不稳定

5. 推荐刷题清单

题号	题名	难度	关键
215	第 K 个最大元素	Medium	小根堆/快选
347	前 K 个高频元素	Medium	频率+堆
295	数据流中位数	Hard	双堆
23	合并 K 个升序链表	Hard	小根堆

延伸阅读

• LeetCode 官方
• 代码随想录
• 《算法导论》(CLRS) 第 6 章