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