LeetCode 347: 前 K 個高頻元素 Top K Frequent Elements

題目:

給定一個非空的整數數組，返回其中出現頻率前 K 高的元素。

Given a non-empty array of integers, return the K most frequent elements.

示例 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 是數組的大小。

Note:

• You may assume k is always valid, 1 ≤ k ≤ number of unique elements.
• Your algorithm's time complexity must be better than O(n log n), where n is the array's size.

解題思路:

​ 這道題大致解題步驟是: 頻率統計 --> 按頻率排序 --> 返回頻率最高的前 K 個元素

注意點:

• 題目要求時間複雜度優於 O(n log n)

首先頻率統計最優雅的方法應該是藉助哈希映射, key 爲元素, value 爲頻率. 其時間複雜度爲 O(n)

重點是返回前 K 個頻率最高的元素, 所以另一種更簡單的方法是直接藉助 堆(優先隊列) 這種數據結構

維護一個 大小爲 K 的堆來動態存儲前 K 個頻率最高的元素, 其時間複雜度爲 O(n)

代碼:

Java:

class Solution {
    public List<Integer> topKFrequent(int[] nums, int k) {
        // 建立哈希映射
        HashMap<Integer, Integer> count = new HashMap();
        // 頻率統計
        for (int n : nums) count.put(n, count.getOrDefault(n, 0) + 1);

        // 建立優先隊列, 藉助 Lambda 表達式
        PriorityQueue<Integer> heap = new PriorityQueue<Integer>((a, b) -> count.get(a) - count.get(b));
        // 也可以藉助 compare 比較函數
        // PriorityQueue<Integer> heap = new PriorityQueue<>(new Comparator<Integer>() {
        //     @Override
        //     public int compare(Integer a, Integer b) {
        //         return map.get(a) - map.get(b);
        //     }
        // });
        
        // 維護一個大小爲 k 的已排序的優先隊列
        for (int n : count.keySet()) {
            heap.add(n);
            if (heap.size() > k)
                heap.poll();
        }

        // 返回結果
        List<Integer> top_k = new LinkedList();
        while (!heap.isEmpty())
            top_k.add(heap.poll());
        return top_k;
    }
}

Python:

Python 基礎庫裏的 heapq 堆數據結構, 有兩個函數:

• nlargest
• nsmallest

例如

heapq.nsmallest(n, nums)

表示取迭代器 nums 前 n 個最大元素, 該函數還能接受一個 key 關鍵字,以應對複雜的數據結構

結合 collections.Counter() 頻率統計函數, 兩行代碼即可解決

class Solution:
    def topKFrequent(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: List[int]
        """ 
        count = collections.Counter(nums)   
        return heapq.nlargest(k, count.keys(), key=count.get) 

注意體會關鍵字參數的作用: key=count.get

歡迎關注微.信公.衆號: 愛寫Bug