Top K Frequent Elements

Count the occurrences first, then select the k entries with the highest counts. The interesting part is the selection step, and there are two good ways to do it.

The heap approach keeps a min heap of size k of the entries seen so far. As you walk the counts, if the current entry has a higher count than the smallest in the heap, replace that smallest entry. At the end the heap holds the k most frequent items. The bucket approach groups the entries by count, so you build a list of buckets where bucket number j holds the items that appear j times, and then you walk the buckets from the highest count down collecting items until you have k of them.

The mechanism is that once you have the counts, the problem becomes selecting the k largest values of a small collection, and both approaches exploit that structure. The heap avoids sorting the whole count table, and the bucket sort exploits the fact that the counts are bounded by the length of the input.

The trade-off to name is that the heap approach is O of n log k time, where n is the number of items and k the number you want, while the bucket approach is O of n time because the buckets are bounded by the maximum frequency. The bucket approach does need that extra array of buckets, which is proportional to the maximum count. If k is close to the number of distinct items, sorting the counts is also perfectly reasonable and is the simplest of the three to write.

Count the occurrences, then keep the k entries with the highest counts using the heapq helpers.

from collections import Counter
import heapq


def top_k_frequent(values, k):
    counts = Counter(values)
    return heapq.nlargest(k, counts, key=counts.get)


data = [1, 1, 1, 2, 2, 3]
print(top_k_frequent(data, 2))
print(top_k_frequent(data, 3))

The outputs are 1 and 2 and then 1, 2, 3. The value 1 appears three times, 2 appears twice and 3 appears once, so they come out in that order.

The mechanism is that nlargest keeps a heap of size k internally, comparing entries using the key function you supply. Passing counts.get means the comparison uses the occurrence count rather than the value itself, which is the part that makes this correct.

The trade-off to name is that the heap approach never sorts the whole count table, so it costs O of n log k where n is the number of distinct items. If you need the result in a specific order, for example the most frequent first, you have to sort the selected entries afterwards, and that adds a step. The library helper also builds the entire count table before selecting anything, so for an enormous input where you only want the top few you would want to count and select in a single streaming pass instead.