java - 前 K 个频繁元素

Question

我正在处理一个面试问题，我需要在给定非空整数数组的情况下返回 k 个最频繁出现的元素。

Input: nums = [1,1,1,2,2,3], k = 2
Output: [1,2]

这是我的代码:

  public static void main(String[] args) {
    System.out.println(topKFrequent(new int[] {1, 1, 1, 2, 2, 3, 3}, 1));
  }

  private static List<Integer> topKFrequent(int[] nums, int k) {
    // freq map
    Map<Integer, Integer> freq = new HashMap<Integer, Integer>();
    for (int n : nums) {
      freq.put(n, freq.getOrDefault(n, 0) + 1);
    }
    // bucket sort on freq
    List<Integer>[] bucket = new LinkedList[nums.length + 1];
    for (int i = 0; i < bucket.length; i++)
      bucket[i] = new LinkedList<>();
    for (int key : freq.keySet()) {
      bucket[freq.get(key)].add(key);
    }
    // gather result
    List<Integer> res = new ArrayList<>();
    for (int i = bucket.length - 1; i >= 0; i--) {
      res.addAll(bucket[i]);
      if (res.size() >= k)
        break;
    }
    return res;
  }

我想检查一下这段代码是否有任何改进/优化？如果存在多个具有相同频率的数字，我会感到困惑。遇到这种情况怎么办？

Answer 1

这部分

List<Integer>[] bucket = new LinkedList[nums.length + 1];
for (int i = 0; i < bucket.length; i++)
  bucket[i] = new LinkedList<>();
for (int key : freq.keySet()) {
  bucket[freq.get(key)].add(key);
}

我会用一个循环来写

List<Integer>[] bucket = new LinkedList[nums.length + 1];
for (int key : freq.keySet()) {
  List<Integer> currentBucket = bucket[freq.get(key)];
  if (currentBucket == null) {
      currentBucket = new LinkedList<Integer>();
      bucket[freq.get(key)] = currentBucket;
  }
  currentBucket.add(key);
}

在这种情况下是最后一部分

List<Integer> res = new ArrayList<>();
for (int i = bucket.length - 1; i >= 0; i--) {
  res.addAll(bucket[i]);
  if (res.size() >= k)
    break;
}

需要跳过空值

List<Integer> res = new ArrayList<>();
for (int i = bucket.length - 1; i >= 0; i--) {
  if (bucket[i] == null) {
      continue;
  }
  res.addAll(bucket[i]);
  if (res.size() >= k)
    break;
}