Top K frequent elements I am working on an interview question where I need to return the k most frequent elements given a non-empty array of integers.

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

Here is my code:

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

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


I wanted to check whether there is any improvement/optimization can be made in this code?
Instead of looping over the freq.keySet, I recommend using freq.forEach((k, v) -&gt; bucket[v].add(k)). This avoids the map lookups. In your current code, for each key you do a search in the map, which takes time. By using forEach, this additional search is avoided.

Do you actually have any performance requirements? If not, I'd say the code is perfect. It is easy to read and understand, and well documented.

Your code might return more than the k elements that are requested, because of the addAll. This should be discussed with the interviewer whether it is ok.

Too bad that Java doesn't have a built-in Histogram class. Things would have been much simpler that way:

new Histogram(nums).stream()
    .limit(k)
    .forEach(System.out::println);