~/problems / Tries

Basics: count words by prefix with a trie

easy basics ~10 min

Implement class PrefixCounter, a trie that remembers how many words pass through each node:

  • PrefixCounter() starts empty.
  • add(word) stores word (lowercase letters, length >= 1). Adding the same word again stores it again, so it's counted twice.
  • count(prefix) returns how many stored words start with prefix. The empty prefix "" matches every stored word.
p = PrefixCounter()
for w in ["tea", "ten", "to", "tea"]:
    p.add(w)
p.count("te")     # 3   ("tea" twice, "ten")
p.count("tea")    # 2   (a whole word is also a prefix of itself)
p.count("teas")   # 0
p.count("")       # 4

Both methods must cost O(length of the argument), however many words are stored: the tests add 20,000 words and then ask 20,000 questions, so comparing the prefix against every word is too slow.

Use nested dicts as nodes: each node maps a letter to a child node, and also keeps a counter of how many words went through it (for example under a key like "#" that can't be a letter, or use a small Node class).

Show hint

in add, walk down from the root creating missing children with setdefault and add 1 to the count of every node you step into; count walks the same path and returns the count at the last node, or 0 as soon as a letter is missing.

Topic: Trie. Prefix tree; longest-match tokenizing.

Read the visual guide
0:00
Ctrl ' run · Ctrl ↵ submit
esc