Top K frequent elements

Medium TimeO(n) SpaceO(n)

Given an array nums and a number k, return the k values that appear most often in it. The array is unsorted, k is never larger than the number of distinct values, and the answer is those values themselves rather than their counts.

Examples

Example 1

Input
nums = [1, 1, 1, 2, 2, 3]k = 2
Output
[1,2]
1 appears three times and 2 appears twice, so those are the top two.

Example 2

Input
nums = [5, 5, 7]k = 1
Output
[5]
5 appears twice and 7 once, and only one value is asked for.

The Code

function topKFrequent(nums, k) {
  const counts = new Map();
  for (let i = 0; i < nums.length; i++) {
    counts.set(nums[i], (counts.get(nums[i]) || 0) + 1);
  }
  const buckets = [];
  counts.forEach((count, value) => {
    if (buckets[count] === undefined) buckets[count] = [];
    buckets[count].push(value);
  });
  const out = [];
  for (let count = nums.length; count >= 1 && out.length < k; count--) {
    if (buckets[count] === undefined) continue;
    for (let i = 0; i < buckets[count].length && out.length < k; i++) {
      out.push(buckets[count][i]);
    }
  }
  return out;
}
topKFrequent([1, 1, 1, 2, 2, 3], 2);
Done

The first 12 calls, of 18. This one does not fit on a page.

Step through topKFrequent([1, 1, 1, 2, 2, 3], 2) call by call

Explanation

Sorting the values by how often they appear would answer this, and costs O(n log n). But a count is a small whole number — at most the length of the array — so it can be used as an index: a value seen three times is pushed into buckets[3]. Walking those buckets from the highest count down then yields the values in frequency order without anything ever being sorted, and the walk stops as soon as out holds k.

  1. 1

    counts starts empty. At i = 0 the value 1 is seen for the first time, so counts.set(1, 1).

    i
    1
    0
    1
    1
    1
    2
    2
    3
    2
    4
    3
    5

    counts={}

    counts → {1: 1}

  2. 2

    Two more 1s follow, each adding one to the same entry.

    1
    0
    1
    1
    i
    1
    2
    2
    3
    2
    4
    3
    5

    counts={1: 2}

    counts → {1: 3}

  3. 3

    The 2s take an entry of their own, by exactly the same line.

    1
    0
    1
    1
    1
    2
    2
    3
    i
    2
    4
    3
    5

    counts={1: 3, 2: 1}

    counts → {1: 3, 2: 2}

  4. 4

    The last value, the 3, gets its own entry and the counting pass ends. One pass, and nothing ordered.

    1
    0
    1
    1
    1
    2
    2
    3
    2
    4
    i
    3
    5

    counts={1: 3, 2: 2}

    counts → {1: 3, 2: 2, 3: 1}

  5. 5

    Now each value is filed under its own count: the 1 was seen three times, so it is pushed into buckets[3]; the 2 into buckets[2]; the 3 into buckets[1].

    1
    0
    1
    1
    1
    2
    2
    3
    2
    4
    i
    3
    5

    counts={1: 3, 2: 2, 3: 1}buckets=[]

    buckets → 3: [1] · 2: [2] · 1: [3]

  6. 6

    The last loop counts count down from 6, skipping the buckets that are empty, and pushes what it finds into out until out holds k values. Bucket 3 gives the 1 and bucket 2 gives the 2 — that is two, so the loop stops there and the 3 in bucket 1 is never reached.

    1
    0
    2
    1

    out=[]k=2

    return [1, 2]