DSA Top K Frequent Words

The problem

Top K Frequent Words lists the k most frequent words in a text. Frequency is the count of one word inside the text. Words are separated by single spaces in this article.

Equal counts are ordered alphabetically. Alphabetical order compares words letter by letter from the left. That tie rule is part of the answer, not a detail.

cases = [
    ("the day is sunny the the sunny is sunny day", 1, [("sunny", 3)]),
    ("the day is sunny the the sunny is sunny day", 2, [("sunny", 3), ("the", 3)]),
    ("the day is sunny the the sunny is sunny day", 4, [("sunny", 3), ("the", 3), ("day", 2), ("is", 2)]),
    ("moon sun moon sun star", 2, [("moon", 2), ("sun", 2)]),
]
for text, k, expected in cases:
    print("k =", k, "->", expected)

Each pair shows a word and its count. In the main text, sunny and the tie at three. Sunny takes the first place because it comes earlier in alphabetical order.

Intuition: first-seen order is the wrong tie-break

The first idea counts the words and reports the most common ones in the order they were met. Counting is easy. The trap hides in the ties, where first-seen order quietly replaces the stated alphabetical rule.

The text below meets sun before moon. Both appear twice. First-seen order answers sun, and the stated rule demands moon.

from collections import Counter

text = "sun moon sun moon"
counts = Counter(text.split())
print("counts:", sorted(counts.items()))
print("first-seen winner:", counts.most_common(1))
print("rule winner: moon, earlier alphabetically")

A tie is not a rare corner here. Short texts tie constantly. Any method that cannot state its tie order cannot be trusted with the answer.

Brute force: count everything, sort everything

The brute force counts every word once. It then sorts the whole vocabulary by count downward and word upward. The first k entries of that sorted list are the answer.

Sorting by a pair works because the pair compares field by field. A negated count puts larger counts first. The word itself then orders equal counts alphabetically.

from collections import Counter

def top_k_brute(text, k):
    counts = Counter(text.split())
    entries = [(-count, word) for word, count in counts.items()]
    entries.sort()
    return [(word, -neg) for neg, word in entries[:k]], len(entries)

text = "the day is sunny the the sunny is sunny day"
result, vocab = top_k_brute(text, 2)
print("top 2:", result)
print("words sorted in full:", vocab)
result, vocab = top_k_brute("moon sun moon sun star", 2)
print("tie case:", result)

The answers are right. The waste is the full sort of every distinct word, when only k words will ever be read. A vocabulary of thousands pays to order words nobody asked about.

The optimal solution: heapify the entries, pop only k

The method in one sentence: build one heap entry of (-count, word) per distinct word, heapify the entries in one pass, then pop exactly k times. Heapify is the linear-time build that turns a list into a heap in place.

A min-heap pops the smallest entry first. The smallest negated count belongs to the largest real count. Equal negated counts fall through to the word, which pops in alphabetical order for free.

import heapq
from collections import Counter

def top_k(text, k):
    counts = Counter(text.split())
    entries = [(-count, word) for word, count in counts.items()]
    heapq.heapify(entries)
    result = []
    pops = 0
    for _ in range(min(k, len(entries))):
        neg, word = heapq.heappop(entries)
        pops += 1
        result.append((word, -neg))
    return result, pops

cases = [
    ("the day is sunny the the sunny is sunny day", 1, [("sunny", 3)]),
    ("the day is sunny the the sunny is sunny day", 2, [("sunny", 3), ("the", 3)]),
    ("the day is sunny the the sunny is sunny day", 4, [("sunny", 3), ("the", 3), ("day", 2), ("is", 2)]),
    ("moon sun moon sun star", 2, [("moon", 2), ("sun", 2)]),
]
for text, k, expected in cases:
    result, pops = top_k(text, k)
    print("k =", k, "->", result, "match:", result == expected, "pops:", pops)

All four cases pass. Only k pops are paid, whatever the vocabulary size. The tie case answers moon before sun, exactly as the rule states.

Walkthrough: the main text with k = 3

The walk starts from the counted vocabulary. The state is the heap of entries, printed sorted for reading. Each pop removes the current winner and the next winner surfaces.

import heapq
from collections import Counter

text = "the day is sunny the the sunny is sunny day"
counts = Counter(text.split())
print("counts:", sorted(counts.items()))
entries = [(-count, word) for word, count in counts.items()]
heapq.heapify(entries)
print("heapified, sorted for reading:", sorted(entries))
result = []
for step in range(3):
    neg, word = heapq.heappop(entries)
    result.append(word)
    print("pop", step + 1, ":", word, "count", -neg, "- heap left:", sorted(entries))
print("top 3:", result)

Sunny pops before the at equal counts. No sorting step made that happen. The pair ordering inside each entry carried the tie rule through the heap on its own.

Complexity

Counting costs one step per word in the text, so it is O(n) for n words, and both methods pay it. The difference comes after: sorting the vocabulary of V distinct words costs O(V log V), while heapify costs O(V) and each of the k pops costs about log2 V, for O(V + k log V) after counting.

Both methods store one entry per distinct word, so both use O(V) extra space. The heap version saves time, not memory.

for v in (4, 64, 1024):
    levels = v.bit_length()
    print("vocabulary", v, "- sort budget about:", v * levels, " heapify plus 2 pops about:", v + 2 * levels)

Counts beat clocks here. At a vocabulary of 1024, the sort budget runs past ten times the vocabulary. The heap build stays near the vocabulary itself, and the two pops barely register.

Edge cases and pitfalls

A k of zero asks for nothing and answers an empty list. A k beyond the vocabulary answers every word, in rule order, rather than failing. An empty text has no words and answers an empty list.

Case and punctuation are part of a word here. The word Sun differs from sun, and day with a full stop differs from day. Splitting on spaces does no cleaning, and the run below shows both effects.

import heapq
from collections import Counter

def top_k(text, k):
    counts = Counter(text.split())
    entries = [(-count, word) for word, count in counts.items()]
    heapq.heapify(entries)
    result = []
    for _ in range(min(k, len(entries))):
        neg, word = heapq.heappop(entries)
        result.append((word, -neg))
    return result

print("k = 0:", top_k("a a b", 0))
print("k beyond vocabulary:", top_k("b a a", 9))
print("empty text:", top_k("", 3))
print("all tied:", top_k("pear fig kiwi", 3))
print("case counts separately:", top_k("Sun sun Sun", 2))
print("punctuation sticks:", top_k("day day. day", 2))

The all-tied row returns fig, kiwi and pear in alphabetical order. With every count equal, the tie rule becomes the whole ordering, and the entries still carry it.