[LeetCode in Python] 347 (M) top k frequent elements 前 K 个高频元素

题目

https://leetcode-cn.com/problems/top-k-frequent-elements/

给定一个非空的整数数组，返回其中出现频率前k高的元素。

示例 1:

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

示例 2:

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

说明：

你可以假设给定的k总是合理的，且 1 ≤ k ≤ 数组中不相同的元素的个数。
你的算法的时间复杂度必须优于 O(n log n) ,n是数组的大小。

解题思路

• bucket sorting 桶排序
• 先扫描数组，获得数字频率表
• 再构造{频率:[数字]}这样的字典，这就是所谓的桶
• 在上一步，获得频率的最大值max_f
• 从max_f到0逆序循环，按照题意不需要考虑异常情况，整桶添加到结果列表中
• 直到结果列表长度满足k即可退出

代码

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        
        # - statistic frequency
        freq_dict = defaultdict(int)
        for n in nums:
            freq_dict[n] += 1
            
        # - assign buckets, format is {freq:[number]}
        bucket_dict = defaultdict(list)
        max_f = 0
        for n,f in freq_dict.items():
            bucket_dict[f].append(n)
            max_f = max(max_f, f)
        
        # - reverse traverse
        res = []
        for f in range(max_f,0,-1):
            if f in bucket_dict:
                res += bucket_dict[f]
                if len(res) >= k:
                    return res