Reference: LeetCode 347 / LeetCode 692
Difficulty: Medium

Top K Frequent Elements

Given a non-empty array of integers, return the k most frequent elements.

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.

Example: 

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

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

Heap

  • Build a hash map val -> count. $O(N)$
  • Build a min heap of size $k$. $O(N\log{k})$
  • Build an output list. $O(k\log{k})$

Note:

  • If we use max heap, the runtime would be $O(N\log{N})$, since we cannot keep the heap size as $k$. Therefore, for largest problems we use min heap.
  • The resulting order is not required.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
public List<Integer> topKFrequent(int[] nums, int k) {
  // assume nums is non-empty; k is always valid
  
  // count
  Map<Integer, Integer> map = new HashMap<>();
  for (int val : nums) {
    map.put(val, map.getOrDefault(val, 0) + 1);
  }
  
  // priority queue
  List<Integer> result = new ArrayList<>();
  PriorityQueue<Integer> minPQ = new PriorityQueue<>((v1, v2) -> {
    return map.get(v1) - map.get(v2);
  });
  
  for (int val : map.keySet()) {
    minPQ.offer(val);
    if (minPQ.size() > k) {
      minPQ.poll();
    }
  }
  
  // output
  while (minPQ.size() > 0) {
    result.add(minPQ.poll());
  }
  
  return result;
}

Time: $O(N\log{k})$
Space: $O(N)$

Top K Frequent Words

Given a non-empty list of words, return the k most frequent elements.

Your answer should be sorted by frequency from highest to lowest. If two words have the same frequency, then the word with the lower alphabetical order comes first.

Example:

1
2
Input: ["i", "love", "leetcode", "i", "love", "coding"], k = 2
Output: ["i", "love"]

Sorting

For comparator, from left to right: w1, w2 and count(w2), count(w1).

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
public List<String> topKFrequent(String[] words, int k) {
  // assume words is non-empty and k is valid

  // count
  Map<String, Integer> map = new HashMap<>();
  for (String s : words) {
    map.put(s, map.getOrDefault(s, 0) + 1);
  }

  // sorting (descending)
  List<String> result = new ArrayList<>(map.keySet());
  Collections.sort(result, (w1, w2) -> {
    int cmp = map.get(w2) - map.get(w1);
    if (cmp == 0) {
      return w1.compareTo(w2);
    } else {
      return cmp;
    }
  });

  return result.subList(0, k);
}

Time: $O(N\log{N})$
Space: $O(N)$

Heap

Note: Be careful of the code in the comparator.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
public List<String> topKFrequent(String[] words, int k) {
  // assume words is non-empty and k is valid

  // count
  Map<String, Integer> map = new HashMap<>();
  for (String s : words) {
    map.put(s, map.getOrDefault(s, 0) + 1);
  }

  // priority queue
  PriorityQueue<String> minPQ = new PriorityQueue<>((s1, s2) -> {
    int cmp = map.get(s1) - map.get(s2);
    if (cmp == 0) {
      return s2.compareTo(s1); // "z" goes first (alphabet)
    } else {
      return cmp;
    }
  });
  
  LinkedList<String> result = new LinkedList<>();
  
  for (String s : map.keySet()) {
    minPQ.offer(s);
    if (minPQ.size() > k) {
      minPQ.poll();
    }
  }

  // output (ordering matters)
  while (minPQ.size() > 0) {
    result.addFirst(minPQ.poll());
  }
  
  return result;
}

Time: $O(N\log{k})$
Space: $O(N)$