LeetCode 347: 前 K 個高頻元素 Top K Frequent Elements
- 2019 年 12 月 18 日
- 筆記
題目:
給定一個非空的整數數組，返回其中出現頻率前 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
