高效求解前K个高频元素的两种算法

1. 问题描述

给定一个整数数组 nums 和一个整数 k，请返回其中出现频率前 k 高的元素。可以按任意顺序返回答案。

示例 1:

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

示例 2:

code复制输入: nums = [1], k = 1
输出: [1]

提示：

• 1 <= nums.length <= 10^5
• k 的取值范围是 [1, 数组中不相同的元素的个数]
• 题目数据保证答案唯一，换句话说，数组中前 k 个高频元素的集合是唯一的

2. 解题思路

方法一：哈希表 + 堆

1. 统计频率：首先使用哈希表统计每个数字出现的频率。
2. 构建最小堆：维护一个大小为 k 的最小堆，堆中存储的是数字及其频率。遍历哈希表中的每个数字，将其插入堆中。如果堆的大小超过 k，则弹出堆顶元素（频率最小的元素）。
3. 获取结果：最后堆中剩下的元素就是前 k 个高频元素。

时间复杂度：

• 统计频率：O(n)
• 构建堆：O(n log k)
• 总体：O(n log k)

空间复杂度：O(n)

方法二：桶排序

1. 统计频率：同样先使用哈希表统计每个数字出现的频率。
2. 桶排序：将数字按照频率放入对应的桶中，桶的下标表示频率。
3. 从高频率到低频率遍历桶，收集前 k 个高频数字。

时间复杂度：

• 统计频率：O(n)
• 桶排序：O(n)
• 总体：O(n)

空间复杂度：O(n)

3. 代码实现

方法一：哈希表 + 堆

python复制import heapq
from collections import defaultdict

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        # 统计频率
        frequency = defaultdict(int)
        for num in nums:
            frequency[num] += 1
        
        # 构建最小堆
        heap = []
        for num, count in frequency.items():
            heapq.heappush(heap, (count, num))
            if len(heap) > k:
                heapq.heappop(heap)
        
        # 提取结果
        return [num for count, num in heap]

方法二：桶排序

python复制from collections import defaultdict

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        # 统计频率
        frequency = defaultdict(int)
        for num in nums:
            frequency[num] += 1
        
        # 桶排序
        buckets = [[] for _ in range(len(nums) + 1)]
        for num, count in frequency.items():
            buckets[count].append(num)
        
        # 从高频率到低频率遍历桶
        res = []
        for i in range(len(buckets) - 1, -1, -1):
            res.extend(buckets[i])
            if len(res) >= k:
                break
        return res[:k]

4. 复杂度分析

方法一：哈希表 + 堆

• 时间复杂度：O(n log k)，因为每个元素插入堆的时间复杂度为 O(log k)，共有 n 个元素。
• 空间复杂度：O(n)，用于存储哈希表和堆。

方法二：桶排序

• 时间复杂度：O(n)，因为统计频率和桶排序都是线性时间。
• 空间复杂度：O(n)，用于存储哈希表和桶。

5. 测试案例

示例 1:

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

示例 2:

code复制输入: nums = [1], k = 1
输出: [1]

6. 总结

• 方法一适用于 k 较小的情况，因为堆的大小保持在 k，插入和弹出的操作相对高效。
• 方法二在 k 较大时更为高效，因为其时间复杂度是线性的，且实现简单。