347. 前 K 个高频元素

import collections
from typing import List

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        """
        找出 nums 数组中出现频率最高的前 k 个元素 (使用桶排序思想)。

        Args:
            nums: 整数数组。
            k: 需要返回的高频元素的数量。

        Returns:
            一个包含前 k 个高频元素的列表。
        """
        
        # 1. 统计频率
        # freq_map 的形式是 {元素: 频率}
        # 时间复杂度: O(N)
        freq_map = collections.Counter(nums)

        # 2. 创建桶数组
        # 桶的索引代表频率，范围从 0 到 len(nums)
        # buckets[i] 存储所有频率为 i 的元素列表
        # 时间复杂度: O(N) (初始化列表)
        # 空间复杂度: O(N) (最坏情况)
        n = len(nums)
        buckets = [[] for _ in range(n + 1)] 

        # 3. 元素入桶
        # 遍历频率映射，将元素放入对应频率的桶中
        # 时间复杂度: O(N')，N' 是不同元素的数量 (N' <= N)
        for num, freq in freq_map.items():
            buckets[freq].append(num)

        # 4. 按频率逆序收集结果
        result = []
        # 从最大可能的频率 n 开始，向下遍历到频率 1
        # 时间复杂度: O(N) (最坏情况遍历所有桶和元素)
        for i in range(n, 0, -1):
            # 如果当前频率 i 的桶不为空
            if buckets[i]:
                # 将该桶中的所有元素添加到结果中
                for num_in_bucket in buckets[i]:
                    result.append(num_in_bucket)
                    # 如果结果列表已达到 k 个，可以直接返回
                    if len(result) == k:
                        return result
        
        # 如果 unique 元素少于 k 个 (理论上题目保证 k 有效)
        return result

# 示例用法:
# solution = Solution()
# nums1 = [1,1,1,2,2,3]
# k1 = 2
# print(solution.topKFrequent(nums1, k1)) # 输出: [1, 2] (顺序可能不同)

# nums2 = [1]
# k2 = 1
# print(solution.topKFrequent(nums2, k2)) # 输出: [1]