Top K frequent words in Python: heap versus bucket sort

Why this matters

Search boxes, chat apps and analytics pages often show a "top 10" list: the 10 most searched terms, the 10 most used tags, the 10 most common errors in a log. To build such a list, a program needs two steps. It must count how often each item appears. Then it must pick the items with the largest counts.

The counting step is simple with a dict. The picking step has more than one good answer, and the answers differ in speed when the list is long. This lesson builds two of them, a heap and buckets, and compares them with a plain sort.

A dict is a collection that stores pairs of a key and a value, and finds a value from its key quickly. A heap is a structure that gives you the smallest item quickly, without keeping the whole collection sorted. Python's heapq module provides a heap. It is a min-heap, which means the smallest item is always at the top. A bucket is a list that holds all the items with one particular count.

The problem

You get a list of words and a number K. You must return the K words that appear most often. If two words have the same count, the one that comes first in alphabetical order goes first.

Here is an example with K equal to 2.

words:  i, love, leetcode, i, love, coding
counts: i=2, love=2, leetcode=1, coding=1
answer: [i, love]

i and love both appear twice, so they are on top. i comes before love alphabetically, so it is first in the answer. leetcode and coding each appear once and do not make the list.

The work splits into two steps. Step one counts the words. Step two picks the K best. Each step gets its own section below. The code uses these names:

• n is the total number of words in the input list.
• u is the number of different words (the number of distinct words). Because u can never be larger than n, u is at most n.
• K is the number you were asked for.

Step one: count with Counter

collections.Counter is a dict subclass made for counting. Give it any iterable, and it stores each item as a key with its count as the value. Counting costs about n steps in total, because each word is looked at once, and each dict lookup takes about the same time no matter how many words are stored.

{'coding': 1, 'i': 2, 'leetcode': 1, 'love': 2}

The dict is sorted in the output only so that it is easy to read. A Counter does not promise any order when you loop over it. That is fine here, because step two sorts the results by count and by word.

Step two, first method: sort everything

The simplest way to pick the K best is to sort all the distinct words by count, highest first, and take the first K. For words with equal counts, sort them alphabetically.

The sorted function takes a key function. The key here is a lambda (a short unnamed function) that returns a tuple: the negative count, then the word. Sorting by a tuple compares the first part, and uses the second part only when the first parts are equal. The negative count makes the largest count come first.

['i', 'love']

The sort compares pairs of items, and sorting u items takes about u times log u comparisons. The word log here means the base-2 logarithm. You can think of it as "how many times you can split the number in half before reaching 1." For u equal to 1,000,000, log u is about 20. The sort is correct and simple, and it is the reference answer that the other two methods are checked against.

The sort looks at every distinct word, even when K is only 2. That is the waste the next two methods try to avoid.

Step two, second method: a heap of size K

Keep a heap that holds at most K words. The heap is ordered so that its top item is the weakest word kept so far. For each new word, add it. If the heap now holds more than K words, remove the weakest one. At the end, the heap holds the K best words.

Python's heapq is a min-heap, so the top item is the smallest tuple. Each entry is a tuple (count, tie, word). A smaller count means a weaker word, so it sits at the top. For words with equal counts, the weaker one is the one that comes later in alphabetical order. A min-heap pops the smaller tuple first, so the tie part must be smaller for the later word. The code uses the word's position in sorted order, negated, for that part.

Here is the trace for four words with counts the=4, is=3, sunny=2 and day=1, with K equal to 2. The words are handled in the order day, is, sunny, the.

add day (count 1)     heap: day
add is (count 3)      heap: day, is
add sunny (count 2)   heap has 3 words, remove the weakest: day
                      heap: sunny, is            (sunny is now the weakest)
add the (count 4)     heap has 3 words, remove the weakest: sunny
                      heap: is, the

At the end the heap holds is and the. Popping the heap gives the weakest first, so the list comes out as is, the. Reverse it to get the, is, which is the answer in the right order.

['the', 'is']

heappush adds an item to the heap, and heappop removes the top item, which is the weakest one kept. Each push and each pop costs about log K steps, because the heap never holds more than K items. The loop over all u distinct words therefore costs about u times log K steps.

The heap is never larger than K, so each step is cheap when K is small. This is the main advantage of the heap. When K is 10 and there are 1,000,000 distinct words, the sort does about 20 steps per word, while the heap does about 3 steps per word, because log 10 is about 3.3.

The tuple (count, -order, word) is compared from left to right. The word at the end is only a label. The order part decides every tie, because two different words always have different positions.

Step two, third method: buckets by count

The count of a word is a whole number from 1 to n. A word cannot appear more than n times, because there are only n words in the input. So you can make one bucket for each possible count and put each word into the bucket for its count. Then walk from the highest count down to 1 and collect words until you have K of them.

['the', 'is', 'sunny', 'day']

The code sorts the distinct words once, before it fills the buckets. Each bucket then holds its words in alphabetical order, so the walk gives the tie rule for free.

The trace for the same four words, with K equal to 2:

bucket 1: day
bucket 2: sunny
bucket 3: is
bucket 4: the
walk from bucket 4 down: the (1 word so far), is (2 words so far), stop
answer: the, is

The walk stops as soon as it has K words, so it never visits the low buckets when K is small.

Here is the cost. Counting takes about n steps. Filling the buckets takes about u steps. The walk takes about n steps at most. The alphabetical sort of distinct words takes about u times log u steps. The total is about n plus u log u.

The sort is the part that costs. If the problem allowed ties to come out in any order, the sort could go, and the total would be about n plus u, which is faster than the other two methods. The tie rule in this lesson requires the sort, so the bucket method here has the same main cost as the plain sort. Keep that in mind when you read the comparison below.

Comparing the three methods

The table compares the three methods on the picking step. Counting costs about n steps in all three cases, so it is left out.

When K is small compared with u, the heap wins. It never holds more than K words, so most of the work is cheap. When K is close to u, the heap gives little gain, and the plain sort is simpler to read.

The bucket method uses space for n + 1 lists, which can be large if n is large, even when only a few words are distinct. It is the right tool when counts are small and the tie order does not matter.

In practice, pick the heap when you need the top few items from a large set of words. Pick the sort when the list is short, or when the code clarity matters more than speed. Use buckets when counts are small, and you can accept ties in any order.

For a quick one-line answer in real code, heapq.nlargest(k, counts.items(), key=...) runs the heap method for you. The explicit heap in this lesson shows how it works.

Running all three

The program below runs the three functions on two inputs and then checks them against each other on 2,000 random lists. The random lists are built with a fixed seed, so the result is the same on every run of the same Python version.

count = {'coding': 1, 'i': 2, 'leetcode': 1, 'love': 2}
sort    k=2: ['i', 'love']
heap    k=2: ['i', 'love']
buckets k=2: ['i', 'love']
sort    k=4: ['the', 'is', 'sunny', 'day']
heap    k=4: ['the', 'is', 'sunny', 'day']
buckets k=4: ['the', 'is', 'sunny', 'day']
random trials with mismatch: 0

The three methods agree on every test. The check against the sort is what makes the heap and bucket versions trustworthy: when two methods must give the same answer, a difference shows a bug in one of them.

Common mistakes

Negating the count by mistake

heapq only builds a min-heap, so the item on top is always the smallest. The top-K heap must drop the weakest word, so it stores the plain count. Negating the count is an easy slip. The code runs without error and keeps the wrong words.

The test below uses the four-word example (the, is, sunny, day) with K equal to 2, and pushes the negated count. The top item is now the most frequent word, so the heap drops the best words and keeps the worst.

['day', 'sunny']

Those are the two least frequent words. The fix is to push the plain count, (c, word), so that the weakest word is the smallest item and is the one removed.

Pushing items that cannot be compared

A heap compares its items with <. A dict does not support <, so pushing dicts fails. The error names the types involved:

import heapq
 
heap = []
heapq.heappush(heap, {"a": 1})
heapq.heappush(heap, {"b": 2})

TypeError: '<' not supported between instances of 'dict' and 'dict'

Push tuples whose first parts can be compared, such as (count, word), as the code above does.

Sorting ascending by mistake

sorted sorts from the smallest value to the largest. If you leave out the minus sign or the reverse=True argument, the first entry is the least frequent word. The example below uses the four-word list.

first entry: ('day', 1)

The word day appears once, which is the lowest count. For the top words, use key=lambda kv: -kv[1], or reverse=True with the count as the key.

Sizing the buckets by K instead of by n

The bucket index is a count, and a count can be as large as n. If you make only K + 1 buckets, a word that appears more often than K times will be out of range.

from collections import Counter
 
counts = Counter(["the", "the", "the", "the"])
buckets = [[] for _ in range(3)]
for w, c in counts.items():
    buckets[c].append(w)

IndexError: list index out of range

The list has 3 slots, numbered 0 to 2, and the code asked for slot 4. Make the list n + 1 slots long, as the top_k_buckets function does.

Practice

Exercise 1. Count the letters of the word mississippi, then print the top 2 letters with the sort method. Ties are broken alphabetically.

Show the answer

Counts are i=4, s=4, p=2, m=1. The two letters with count 4 are i and s. They are alphabetical, so the answer is ['i', 's'].

['i', 's']

Exercise 2. Words are a, b, b, c. Find the top 1 word with the heap method. Write down the heap after each step.

Show the answer

Counts are a=1, b=2, c=1. The heap works on the distinct words, one entry per word, so it sees three entries. The heap holds one item. Here is the heap after each entry:

a (count 1)   heap: a
b (count 2)   two items, remove the weaker one, a: heap: b
c (count 1)   two items, remove the weaker one, c (same count, and c comes later): heap: b

b has count 2, which is more than the others, so it is never the weakest and it stays. The answer is ['b'].

Exercise 3. The bucket method uses n + 1 buckets. Explain why u + 1 buckets would be wrong.

Show the answer

Bucket numbers are counts, not the number of distinct words. A word can appear n times. For example, the input a, a, a, a has u equal to 1 and n equal to 4. The word a has count 4, so it needs bucket 4. A list with u + 1 slots, which is 2 slots here, does not have bucket 4, so the code fails with IndexError.

Exercise 4. The bucket method sorts the distinct words before it fills the buckets. Change it so that ties come out in reverse alphabetical order. Which line changes?

Show the answer

Change the line for w in sorted(counts): to for w in sorted(counts, reverse=True):. The buckets are filled in reverse alphabetical order, so the walk returns the tied words from the last letter to the first. The rest of the function stays the same.

Exercise 5. Why does top_k_heap reverse result at the end?

Show the answer

heappop removes the weakest word first, so the list is built from weakest to strongest. The reverse turns it into strongest first, which is the order the problem asks for.

Summary

• Count words with Counter. This step looks at each word once.
• Sorting all distinct words costs about u times log u, and it is the simplest correct way to pick the top K.
• A min-heap of size K keeps only the K best words, and costs about u times log K. It is faster when K is small.
• Buckets by count avoid comparing counts, but the alphabetical tie rule needs a sort of the distinct words, so the cost is the same order as the sort.
• Test the faster methods against the plain sort on random inputs. A mismatch shows a bug.