Swift 实战：用堆结构高效搞定 LeetCode 347 前 K 个高频元素

前言

你有没有遇到过这样一个需求：从一堆数据里找出“出现次数最多的前 K 个元素”？不管是评论区里找最热话题，还是电商数据分析中筛选热门商品，这种“前 K 高频元素”都是非常典型的场景。

LeetCode 347 这道题就完美地还原了这个需求。今天我们用 Swift 带大家写一个实战 Demo，讲清楚“哈希表 + 小顶堆”的高效解法，并且看看它在实际项目中的应用。

描述

题目要求：给你一个整数数组 nums 和一个整数 k，请你返回出现频率前 k 高的元素，返回的顺序不限。

题目示例：

示例 1：

输入: nums = [1,1,1,2,2,3], k = 2
输出: [1,2]

示例 2：

输入: nums = [1], k = 1
输出: [1]

题目约束：

• 1 <= nums.length <= 10^5
• k 一定不会超过数组中的唯一元素个数。
• 答案唯一。

进阶要求：

• 算法时间复杂度必须优于 O(n log n)。

解决方案

这道题本质上是频率统计 + 高效选出前 K 个元素的问题。

解题步骤：

1. 用哈希表统计频率。

2. 用一个小顶堆（优先队列）维护前 K 高频元素。

   • 堆的大小最多为 K。
   • 每次向堆中加入新元素时，如果堆的大小超过 K，就移除堆顶（频率最小的元素）。

3. 最后堆中的元素就是答案。

Swift 没有内置堆（优先队列），不过我们可以借助 Heap 结构自己实现。也可以用数组 + sort 做最简单的模拟，但为了满足“比 O(n log n) 更优”的复杂度，必须用堆。

分析问题并实现代码

我们一步步用 Swift 来实现这个“哈希表 + 小顶堆”的方案。

代码实现：

import Foundation

struct Heap<T> {
    var elements: [T]
    let priorityFunction: (T, T) -> Bool

    init(elements: [T] = [], priorityFunction: @escaping (T, T) -> Bool) {
        self.elements = elements
        self.priorityFunction = priorityFunction
        buildHeap()
    }

    mutating func buildHeap() {
        for index in (0..<(count / 2)).reversed() {
            siftDown(index)
        }
    }

    var isEmpty: Bool { elements.isEmpty }
    var count: Int { elements.count }

    mutating func siftDown(_ index: Int) {
        var parent = index
        while true {
            let left = 2 * parent + 1
            let right = 2 * parent + 2
            var candidate = parent

            if left < count && priorityFunction(elements[left], elements[candidate]) {
                candidate = left
            }
            if right < count && priorityFunction(elements[right], elements[candidate]) {
                candidate = right
            }
            if candidate == parent { return }
            elements.swapAt(parent, candidate)
            parent = candidate
        }
    }

    mutating func siftUp(_ index: Int) {
        var child = index
        var parent = (child - 1) / 2
        while child > 0 && priorityFunction(elements[child], elements[parent]) {
            elements.swapAt(child, parent)
            child = parent
            parent = (child - 1) / 2
        }
    }

    mutating func insert(_ value: T) {
        elements.append(value)
        siftUp(count - 1)
    }

    mutating func remove() -> T? {
        guard !isEmpty else { return nil }
        elements.swapAt(0, count - 1)
        let value = elements.removeLast()
        siftDown(0)
        return value
    }

    func peek() -> T? {
        return elements.first
    }
}

func topKFrequent(_ nums: [Int], _ k: Int) -> [Int] {
    var frequencyMap = [Int: Int]()
    for num in nums {
        frequencyMap[num, default: 0] += 1
    }

    var heap = Heap<(Int, Int)>(priorityFunction: { $0.1 < $1.1 })

    for (num, freq) in frequencyMap {
        heap.insert((num, freq))
        if heap.count > k {
            heap.remove()
        }
    }

    var result = [Int]()
    while !heap.isEmpty {
        if let top = heap.remove() {
            result.append(top.0)
        }
    }
    return result.reversed()
}

代码讲解：

1. Heap 结构：

   • 这是一个通用的最小堆实现，支持插入、删除、堆化。

2. 频率统计：

   • 先用 Dictionary 把每个数字的频率都统计出来。

3. 小顶堆维护前 K 大频率元素：

   • 把频率放入堆中，当堆的元素数量超过 K 时，移除堆顶（频率最小的那个）。

4. 返回结果：

   • 堆中的元素就是前 K 个高频元素，最后再 reverse 一下（因为堆是从小到大排序的）。

示例测试和结果

测试用例 1：

let nums = [1,1,1,2,2,3]
let k = 2
print(topKFrequent(nums, k))  // 输出: [1, 2]

测试用例 2：

let nums2 = [1]
let k2 = 1
print(topKFrequent(nums2, k2))  // 输出: [1]

结果输出：

[1, 2]
[1]

时间复杂度

• 频率统计：O(n)，n 是数组长度。

• 堆插入操作：每次 O(log k)，最多 n 次插入。

• 整体时间复杂度：O(n log k)。

  • 由于 k ≪ n，log k 是个很小的数，所以性能非常高效，满足题目“优于 O(n log n)”的要求。

空间复杂度

• 哈希表空间 O(n)。
• 堆的空间 O(k)。
• 总体空间复杂度是 O(n)。

总结

这道题的解法思路其实非常贴近实际业务场景，比如：

• 评论区的热评排序（前 K 个点赞最多的评论）。
• 实时热门搜索词（前 K 个搜索频率最高的关键词）。
• 电商平台的爆品榜单（前 K 个销量最高的商品）。

通过“哈希表统计 + 小顶堆筛选”，我们用 Swift 高效完成了数据流中的“前 K 高频元素”需求。这种组合拳式的解法，在面试与实战项目中都非常实用，值得掌握。