Top K frequent elements

Question

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<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 ArrayList[nums.length + 1];
    for (int i = 0; i < bucket.length; i++)
      bucket[i] = new ArrayList<>();
    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;
  }

I wanted to check whether there is any improvement/optimization can be made in this code?

Answer 1

Instead of looping over the freq.keySet, I recommend using freq.forEach((k, v) -> 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);

Answer 2

  private static List<Integer> topKFrequent(int[] nums, int k) {

The specification as described explains why the input uses int[], but doesn't seem to fix the return type. Why List<Integer>? I can see an argument for int[] (consistency with the input), and I can see an argument for Set<Integer> (encodes the constraint that the returned values will be distinct), but I can't see a reason for List<Integer>.

---

    // freq map
    Map<Integer, Integer> freq = new HashMap<Integer, Integer>();
    for (int n : nums) {
      freq.put(n, freq.getOrDefault(n, 0) + 1);
    }

Nice use of the recent API additions.

---

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

Why the inconsistency over whether to use {} for a single statement for body?

The bucket sort is time-efficient, but as implemented it's not very memory-efficient. Have you considered doing a first pass over freq to find the maximum frequency, so that bucket can be allocated to the smallest possible size? Have you considered allocating bucket[i] = new ArrayList<>(); on demand, the first time you encounter a frequency of i?

Iterating over a map's keySet() and then calling into the map is inefficient. When HashMap works well it's only a small constant factor inefficiency, but to illustrate the problem suppose that you were using TreeMap. Then keySet() takes \$O(m)\$ time (where \$m\$ is the number of entries in the map), but the \$m\$ calls to get() take a total of \$O(m \lg m)\$ time. Iterating over entrySet() takes \$O(m)\$ time and saves the extra work of looking up the entries by 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;
    }

What about corner cases? If there are fewer than k distinct values in nums, should the method throw an exception or return all of the distinct values? If res.size() > k should you discard some elements from the smallest bucket in order to return only kelements? The specification as quoted is unclear, so maybe you should push for clarification.