Counting words took me twenty lines, and it was slow

See the problem

Your support team wants the three words customers use most in their messages this week. You write a loop that counts each word, then sorts the counts.

On a test file of 50 messages, it’s instant. On the real week, 200,000 words, the script runs for minutes. And the code is twenty lines for something that sounds like one.

Define the task

• Input: a text and a number k.
• Output: the k most common words, as (word, count) pairs, most common first.
• Rules: ignore case. Words with the same count keep the order they first appeared in. Fewer than k words means return them all.
• The hard part: fast on 200,000 words, and short enough to read at a glance.

Example: top_words("the cat and The dog and THE end", 2) returns [("the", 3), ("and", 2)].

Pause and think

⏸ Pause and think

Counting things is so common that the standard library has a dict made for it. Do you know its name?

Give it two minutes with pen and paper before you open a hint.

Hints open one at a time:

Hint 1 of 3

words.count(w) reads every word again. Do it once per word, and the work is words × different words.

Hint 2 of 3

A dict of word → count, filled in one pass, reads each word once.

Hint 3 of 3

collections.Counter(words) is that dict, and .most_common(k) gives the top k.

Understand the idea

Before writing a loop, check whether the standard library has it. These come up most:

Tool	From	What it does
Counter(items)	collections	Counts each item. .most_common(k) gives the top k.
defaultdict(list)	collections	A dict that makes a fresh [] for a missing key.
deque	collections	A list with fast adds and removes at both ends.
chain(a, b)	itertools	Walks a, then b, without joining them.
islice(it, n)	itertools	Takes the first n items of any stream.
groupby(items, key)	itertools	Groups items that sit next to each other.
cache, partial, reduce	functools	Remembers results, fixes some arguments, folds a list into one value.

Counter is a dict underneath, so counting is one pass. most_common(k) uses a heap, which keeps only the best k while it looks.

★Key insight

The standard library already has the loops you keep writing. Counter counts, defaultdict groups, itertools streams, and functools caches. They’re shorter, tested and usually faster.

Climb the ladder

Here n is the number of words, u the number of different words, and k how many you want.

Okay

words.count() for each word

The idea: For each different word, count it with words.count(word), then sort.

Short, and right on small text.

def top_words(text: str, k: int) -> list[tuple[str, int]]:
    words = text.lower().split()
    counts = []
    for word in dict.fromkeys(words):  # each word once, in first-seen order
        counts.append((word, words.count(word)))  # count() reads every word again
    counts.sort(key=lambda pair: pair[1], reverse=True)
    return counts[:k]

Time

O(n × u)

Space

O(u)

Why

count() reads all n words, once for each of the u different words.

Where it breaks: 200,000 words with 20,000 different ones is 4 billion steps. The speed test runs out of time.

Better

A dict, then a sort

The idea: Count every word in one pass with a dict, then sort all the counts and keep k.

The waste in Okay: it reads the whole text again for every word.

Python’s sort is stable, so words with the same count keep their first-seen order.

Highlighted lines are new in this rung.

def top_words(text: str, k: int) -> list[tuple[str, int]]:
    counts: dict[str, int] = {}
    for word in text.lower().split():
        if word in counts:
            counts[word] += 1
        else:
            counts[word] = 1
    ranked = sorted(counts.items(), key=lambda pair: pair[1], reverse=True)
    return ranked[:k]  # sorted every word to keep k of them

Time

O(n + u log u)

Space

O(u)

Why

One pass to count, then a sort of the u counts.

What changed from the previous rung

def top_words(text: str, k: int) -> list[tuple[str, int]]:
    words = text.lower().split()
    counts = []
    for word in dict.fromkeys(words):  # each word once, in first-seen order
        counts.append((word, words.count(word)))  # count() reads every word again
    counts.sort(key=lambda pair: pair[1], reverse=True)
    return counts[:k]
    counts: dict[str, int] = {}
    for word in text.lower().split():
        if word in counts:
            counts[word] += 1
        else:
            counts[word] = 1
    ranked = sorted(counts.items(), key=lambda pair: pair[1], reverse=True)
    return ranked[:k]  # sorted every word to keep k of them

Where it breaks: nine lines for a count, and it sorts all u words to keep three of them.

Best

Counter.most_common

The idea: Counter counts in one pass; most_common(k) keeps the top k with a heap.

The waste in Better: hand-written counting, and a full sort.

One line that any Python reader knows at a glance.

Highlighted lines are new in this rung.

from collections import Counter


def top_words(text: str, k: int) -> list[tuple[str, int]]:
    return Counter(text.lower().split()).most_common(k)

Time

O(n + u log k)

Space

O(u)

Why

One pass to count; the heap holds only k counts while it looks at u.

What changed from the previous rung

from collections import Counter


def top_words(text: str, k: int) -> list[tuple[str, int]]:
    counts: dict[str, int] = {}
    for word in text.lower().split():
        if word in counts:
            counts[word] += 1
        else:
            counts[word] = 1
    ranked = sorted(counts.items(), key=lambda pair: pair[1], reverse=True)
    return ranked[:k]  # sorted every word to keep k of them
    return Counter(text.lower().split()).most_common(k)

When Best isn’t best: for a file too big for memory, count with Counter.update() one chunk at a time. And for real text, split on more than spaces: re.findall(r"[a-z']+", text).

Write it yourself

Write top_words. The last test is 200,000 words, with 20,000 different ones.

The k most common words

Ctrl + Enter runs

from collections import Counter def top_words(text, k): return []

Run testsVisualizeReset codeLoads about 6 MB the first time. Best on Wi-Fi.

1. ○Not run yet: top_words('the cat and The dog and THE end', 2) → [['the', 3], ['and', 2]]
2. ○Not run yet: top_words("b a b a c", 2): ties keep their order → [['b', 2], ['a', 2]]
3. ○Not run yet: top_words('', 3) → []
4. ○Not run yet: top_words(200,000 words, 3) → [['w0', 10], ['w1', 10], ['w2', 10]]

Try your own input

Arguments, separated by commas as in Python

top_words()Run onceVisualize

Compare

Way	Time	Lines	Use it when
Okay: count() per word	O(n × u)	6	Tiny text, once
Better: a dict and a sort	O(n + u log u)	9	You need every count, sorted
Best: Counter.most_common(k)	O(n + u log k)	1	Almost always

Common mistakes and edge cases

!groupby only groups neighbours

itertools.groupby starts a new group whenever the key changes. On ["a", "b", "a"] you get three groups, not two. Sort by the same key first, or use defaultdict(list) instead.

• A missing key in a Counter is 0, not a KeyError. A defaultdict adds the key when you read it, so if d[k]: grows the dict.
• Counter adds up: Counter(a) + Counter(b) sums counts, and drops any that end at 0 or below.
• islice and chain are lazy: they give items one at a time. Wrap them in list() to see them all.
• deque.popleft() is O(1), while list.pop(0) is O(n). Use a deque for queues.
• functools.cache keeps every result forever. Use lru_cache(maxsize=…) when the inputs never repeat much.

Say it in the interview

Questions to ask first

• Is the whole text in memory, or does it arrive as a stream?
• What counts as a word: are “Cat” and “cat,” the same?

A 60-second walkthrough

“I’d lower-case and split, then use collections.Counter, which counts in one pass: O(n). most_common(k) uses a heap to pick the top k in O(u log k), where u is the number of different words. Ties keep first-seen order. A hand-written version is a dict plus a sort of all u words, which is O(u log u) and more code to get wrong. If the text is a stream, I’d call counter.update() on each chunk, so memory stays at one entry per different word.”

Likely follow-ups

• What if u is too big for memory too? Count roughly with a Count-Min Sketch, or split by hash into files and count each.
• defaultdict or setdefault? defaultdict(list) makes the list only when the key is missing; setdefault(k, []) builds a new [] every call.
• How does most_common break ties? Items with equal counts stay in the order they were first counted.

Practice

Three problems, from easy to hard. Each has hints, tests you can run here, and its own Okay, Better and Best answers. Try before you open them.

1. Group words by first letter (easy)

Write by_letter(words), which returns a plain dict from each first letter to the words that start with it, in order. Skip empty strings.

Hint 1 of 3

A defaultdict(list) makes the empty list for you.

Hint 2 of 3

An empty string has no word[0]. Skip it.

Hint 3 of 3

Return dict(groups), so callers get a plain dict.

Group words by first letter

Ctrl + Enter runs

from collections import defaultdict def by_letter(words): return {}

Run testsVisualizeReset codeLoads about 6 MB the first time. Best on Wi-Fi.

1. ○Not run yet: by_letter(['apple', 'bean', 'avocado', 'beet', 'cherry']) → {'a': ['apple', 'avocado'], 'b': ['bean', 'beet'], 'c': ['cherry']}
2. ○Not run yet: by_letter(['', 'a']) → {'a': ['a']}
3. ○Not run yet: Extra test

Extra tests check tricky cases. You see their input only if they fail.

Try your own input

Arguments, separated by commas as in Python

by_letter()Run onceVisualize

Show the Okay, Better and Best answers

Okay

if the key is there

The idea: Append to the list if the letter is a key, else make a new list.

Clear and correct for normal words.

def by_letter(words: list[str]) -> dict[str, list[str]]:
    groups: dict[str, list[str]] = {}
    for word in words:
        if word[0] in groups:
            groups[word[0]].append(word)
        else:
            groups[word[0]] = [word]
    return groups

Time

O(n)

Space

O(n)

Why

One look-up per word.

Where it breaks: six lines for one idea, and ""[0] raises IndexError.

Better

setdefault

The idea: groups.setdefault(letter, []).append(word).

The waste in Okay: the if and else say the same thing twice.

Highlighted lines are new in this rung.

def by_letter(words: list[str]) -> dict[str, list[str]]:
    groups: dict[str, list[str]] = {}
    for word in words:
        groups.setdefault(word[0], []).append(word)  # makes a new [] every time
    return groups

Time

O(n)

Space

O(n)

Why

One look-up per word.

What changed from the previous rung

def by_letter(words: list[str]) -> dict[str, list[str]]:
    groups: dict[str, list[str]] = {}
    for word in words:
        if word[0] in groups:
            groups[word[0]].append(word)
        else:
            groups[word[0]] = [word]
        groups.setdefault(word[0], []).append(word)  # makes a new [] every time
    return groups

Where it breaks: it still fails on an empty string, and builds a new [] on every call, used or not.

Best

defaultdict(list)

The idea: A defaultdict makes the list only for a new letter; skip empty words; return a plain dict.

The waste in Better: a throwaway list per word, and no guard for "".

Highlighted lines are new in this rung.

from collections import defaultdict


def by_letter(words: list[str]) -> dict[str, list[str]]:
    groups: defaultdict[str, list[str]] = defaultdict(list)
    for word in words:
        if word:  # skip empty strings, which have no first letter
            groups[word[0]].append(word)
    return dict(groups)  # a plain dict, so a missing key is an error again

Time

O(n)

Space

O(n)

Why

One look-up per word.

What changed from the previous rung

from collections import defaultdict


def by_letter(words: list[str]) -> dict[str, list[str]]:
    groups: dict[str, list[str]] = {}
    groups: defaultdict[str, list[str]] = defaultdict(list)
    for word in words:
        groups.setdefault(word[0], []).append(word)  # makes a new [] every time
    return groups
        if word:  # skip empty strings, which have no first letter
            groups[word[0]].append(word)
    return dict(groups)  # a plain dict, so a missing key is an error again

When Best isn’t best: if the words are already sorted, itertools.groupby(words, key=lambda w: w[0]) groups them without a dict.

2. Unique, in order (medium)

Write unique_in_order(*lists), which joins any number of lists and drops repeats, keeping each item where it first appeared. It must be fast on 80,000 items.

Hint 1 of 3

item in a_list is O(n). A set or dict looks up in O(1).

Hint 2 of 3

A set forgets the order.

Hint 3 of 3

dict.fromkeys(items) keeps the first place of each key. chain.from_iterable(lists) walks all the lists.

Unique, in order

Ctrl + Enter runs

from itertools import chain def unique_in_order(*lists): return []

Run testsVisualizeReset codeLoads about 6 MB the first time. Best on Wi-Fi.

1. ○Not run yet: unique_in_order([3, 1, 3], [2, 1, 9]) → [3, 1, 2, 9]
2. ○Not run yet: unique_in_order() → []
3. ○Not run yet: four lists of 20,000 that overlap: how many are left? → 50000

Try your own input

Arguments, separated by commas as in Python

unique_in_order()Run onceVisualize

Show the Okay, Better and Best answers

Okay

Check a list

The idea: Append each item if it is not already in the result list.

Right, and keeps the order.

def unique_in_order(*lists: list) -> list:
    seen = []
    for items in lists:
        for item in items:
            if item not in seen:  # a list search: O(n) each time
                seen.append(item)
    return seen

Time

O(n × u)

Space

O(u)

Why

Each in check searches up to u kept items.

Where it breaks: 80,000 items checked against up to 50,000 kept ones is about 2 billion steps.

Better

A set

The idea: list(set(chain(*lists))).

The waste in Okay: a list search where a hash look-up would do.

Highlighted lines are new in this rung.

from itertools import chain


def unique_in_order(*lists: list) -> list:
    return list(set(chain(*lists)))  # a set forgets the order

Time

O(n)

Space

O(u)

Why

Each item is hashed once.

What changed from the previous rung

from itertools import chain


def unique_in_order(*lists: list) -> list:
    seen = []
    for items in lists:
        for item in items:
            if item not in seen:  # a list search: O(n) each time
                seen.append(item)
    return seen
    return list(set(chain(*lists)))  # a set forgets the order

Where it breaks: a set has no order. [3, 1, 3], [2, 1, 9] comes back as [1, 2, 3, 9].

Best

dict.fromkeys

The idea: list(dict.fromkeys(chain.from_iterable(lists))).

The waste in Better: it threw away the order to get speed.

Dict keys are unique, like a set, and keep the order they were added in.

Highlighted lines are new in this rung.

from itertools import chain


def unique_in_order(*lists: list) -> list:
    # dict keys are unique and keep the order they were first added in
    return list(dict.fromkeys(chain.from_iterable(lists)))

Time

O(n)

Space

O(u)

Why

Each item is hashed once.

What changed from the previous rung

from itertools import chain


def unique_in_order(*lists: list) -> list:
    return list(set(chain(*lists)))  # a set forgets the order
    # dict keys are unique and keep the order they were first added in
    return list(dict.fromkeys(chain.from_iterable(lists)))

When Best isn’t best: items must be hashable. For lists of lists, turn each inner list into a tuple first.

3. Chunks of a stream (hard)

Write a generator chunks(items, n) that gives lists of n items, with a shorter last one. It must work on lists, on generators, and on endless streams like itertools.count(). n below 1 raises ValueError("n must be at least 1").

Hint 1 of 3

A generator has no len() and can’t be sliced.

Hint 2 of 3

it = iter(items), then list(islice(it, n)) takes the next n.

Hint 3 of 3

Stop when that list comes back empty.

Chunks of a stream

Ctrl + Enter runs

from itertools import islice def chunks(items, n): yield list(items)

Run testsReset codeLoads about 6 MB the first time. Best on Wi-Fi.

1. ○Not run yet: list(chunks([1, 2, 3, 4, 5], 2)) → [[1, 2], [3, 4], [5]]
2. ○Not run yet: list(chunks(a generator of 0 to 4, 2)) → [[0, 1], [2, 3], [4]]
3. ○Not run yet: list(chunks([1], 0)) → raises ValueError
4. ○Not run yet: Extra test

Extra tests check tricky cases. You see their input only if they fail.

Try your own input

Arguments, separated by commas as in Python

chunks()Run onceVisualize

Show the Okay, Better and Best answers

Okay

Slice by index

The idea: For i in range(0, len(items), n), yield items[i:i + n].

Works on lists and strings.

from collections.abc import Iterator, Sequence


def chunks(items: Sequence, n: int) -> Iterator[list]:
    for i in range(0, len(items), n):  # len() fails on a generator
        yield list(items[i : i + n])

Time

O(n)

Space

O(n)

Why

Each item is copied into one chunk.

Where it breaks: a generator has no len(), so it raises TypeError. And n = 0 gives range’s own error, not ours.

Better

Turn it into a list first

The idea: everything = list(items), then slice that.

The waste in Okay: it only takes things that can be sliced.

Highlighted lines are new in this rung.

from collections.abc import Iterable, Iterator


def chunks(items: Iterable, n: int) -> Iterator[list]:
    everything = list(items)  # never ends on an endless stream
    for i in range(0, len(everything), n):
        yield everything[i : i + n]

Time

O(n)

Space

O(n)

Why

The whole stream is copied into a list first.

What changed from the previous rung

from collections.abc import Iterator, Sequence
from collections.abc import Iterable, Iterator


def chunks(items: Sequence, n: int) -> Iterator[list]:
    for i in range(0, len(items), n):  # len() fails on a generator
        yield list(items[i : i + n])
def chunks(items: Iterable, n: int) -> Iterator[list]:
    everything = list(items)  # never ends on an endless stream
    for i in range(0, len(everything), n):
        yield everything[i : i + n]

Where it breaks: list() reads the whole stream before the first chunk. On an endless stream it never ends; on a big file it fills memory.

Best

islice from one iterator

The idea: Check n, then take list(islice(it, n)) until it comes back empty.

The waste in Better: it holds everything to give out a little at a time.

Highlighted lines are new in this rung.

from collections.abc import Iterable, Iterator
from itertools import islice


def chunks(items: Iterable, n: int) -> Iterator[list]:
    if n < 1:
        raise ValueError("n must be at least 1")
    it = iter(items)
    while chunk := list(islice(it, n)):  # take up to n more, stop when empty
        yield chunk

Time

O(n)

Space

O(chunk)

Why

Only one chunk is held at a time.

What changed from the previous rung

from collections.abc import Iterable, Iterator
from itertools import islice


def chunks(items: Iterable, n: int) -> Iterator[list]:
    everything = list(items)  # never ends on an endless stream
    for i in range(0, len(everything), n):
        yield everything[i : i + n]
    if n < 1:
        raise ValueError("n must be at least 1")
    it = iter(items)
    while chunk := list(islice(it, n)):  # take up to n more, stop when empty
        yield chunk

When Best isn’t best: on Python 3.12 and later, itertools.batched(items, n) does this, giving tuples.

Check yourself

A few quick questions. Ones you miss are remembered, so you can find your weak spots later.

1. Predict the outputWhat does this print?

   from collections import Counter
   
   c = Counter("banana")
   print(c.most_common(2))

   [('a', 3), ('n', 2)][('b', 1), ('a', 3)][('a', 3)]{'a': 3, 'n': 2}

   Check answer
2. Predict the outputWhat does this print?

   from itertools import groupby
   
   print([k for k, _ in groupby("aabaa")])

   ['a', 'b']['a', 'b', 'a']['a', 'a', 'b', 'a', 'a']['b']

   Check answer
3. Multiple choiceYou need a queue that adds at the back and takes from the front, many times. What fits?

   A list with pop(0)collections.deque with popleft()A setA Counter

   Check answer
4. Multiple choiceWhich keeps the first place of each item and drops repeats, in O(n)?

   list(set(xs))list(dict.fromkeys(xs))sorted(set(xs))[x for x in xs if xs.count(x) == 1]

   Check answer

Flashcards

Add these to your review deck. Each card comes back a little later every time you get it right.

• Card 1How do you get the 3 most common items of a list?Show answer
• Card 2What does defaultdict(list) do on a missing key?Show answer
• Card 3What does itertools.islice(it, n) do?Show answer
• Card 4What is the hard part with itertools.groupby?Show answer

Add 4 cards to my review deck

Cheat sheet

• Counter(xs).most_common(k): the top k.
• defaultdict(list): group without if key in d.
• deque: O(1) at both ends.
• chain.from_iterable(lists), islice(it, n): lazy streams.
• groupby groups neighbours only: sort first.
• dict.fromkeys(xs): unique, in order.
• functools.cache, partial, reduce.

Go deeper

• collections.
• itertools, with its recipes section.
• functools.
• Before this: Generators. Next: Type hints.

Can you do these now?

If you’re not sure of any of these, scroll back to that part before moving on.

• Count things with Counter, and get the top k with most_common
• Group things with defaultdict(list), and say how it differs from setdefault
• Chain and slice streams with itertools.chain and islice, without loading them
• Name the functools tools you’ll meet most, like cache, partial and reduce

✓Mark as completeI already know this

Status: Not started