Construct a trie from scratch

The problem, simply

Build a word box with three buttons:

• insert(word) puts a word in the box.
• search(word) says true only if that exact word was put in.
• startsWith(prefix) says true if any word in the box begins with those letters.

All words use the small letters a to z. Below, N is the number of words, L is the length of the word you ask about, and T is the total number of letters inserted.

A story to remember it

Picture it

Examples

Way 1: the simple way

Keep the words in a plain list. To search, compare with every word. To check a start, test every word with startsWith.

1. insert: add the word to the list if it is not there yet.
2. search: walk the list and compare each word.
3. startsWith: walk the list and check whether each word begins with the prefix.

import java.util.*;

class WordList {
    private final List<String> words = new ArrayList<>();

    public void insert(String word) {
        if (!words.contains(word)) words.add(word);
    }

    public boolean search(String word) {
        for (String w : words) {
            if (w.equals(word)) return true;
        }
        return false;
    }

    public boolean startsWith(String prefix) {
        for (String w : words) {
            if (w.startsWith(prefix)) return true;
        }
        return false;
    }
}

How long it takes: every call may compare with all N words, each up to L letters, O(N * L) per call. Extra room: the words themselves, O(N * L). A HashSet makes search fast, but startsWith still has to look at every word.

Way 2: the clever way

Build the hallway house. Each circle, called a node, has 26 slots for children, one per letter, and a flag that says "a word ends here".

1. insert: start at the root. For each letter, if its slot is empty, make a new node there. Step into it. At the end, set the word end flag.
2. search: walk the letters. If a slot is empty, say false. At the end, say whether the word end flag is set.
3. startsWith: the same walk, but at the end say true without checking the flag.

class Trie {
    private static class Node {
        Node[] next = new Node[26];   // one slot per letter
        boolean end;                  // a word ends here
    }

    private final Node root = new Node();

    public void insert(String word) {
        Node cur = root;
        for (char ch : word.toCharArray()) {
            int k = ch - 'a';
            if (cur.next[k] == null) cur.next[k] = new Node();   // grow a branch
            cur = cur.next[k];
        }
        cur.end = true;
    }

    public boolean search(String word) {
        Node node = walk(word);
        return node != null && node.end;   // must reach it and be a word end
    }

    public boolean startsWith(String prefix) {
        return walk(prefix) != null;       // reaching it is enough
    }

    private Node walk(String s) {
        Node cur = root;
        for (char ch : s.toCharArray()) {
            cur = cur.next[ch - 'a'];
            if (cur == null) return null;  // a branch is missing
        }
        return cur;
    }
}

How long it takes: each call walks only the letters of its own word, O(L), no matter how many words are stored. Extra room: at most one node per inserted letter, and each node holds 26 slots, O(26 * T). Shared starts save room, because shared letters are stored once.

Watch out for

• ch - 'a' only works for small letters. For capitals, digits or any Unicode, use a HashMap<Character, Node> for the children.
• Inserting the same word twice must not break anything. The walk reuses every branch and sets the flag again.
• Share the walking code between search and startsWith, so a fix in one helps both.

What interviewers ask next

• Delete a word? Walk to it and clear the flag. Then, going back up, remove nodes that have no children and no flag. A recursive delete does this neatly.
• How many words start with a prefix? Keep a counter in each node and add 1 to it on every node you pass while inserting. The answer is the counter at the prefix node.
• Search with a wildcard . that matches any letter? When you see ., try every child with a small depth first search.
• Why not a hash set of all prefixes? It works for startsWith, but stores every prefix of every word as a separate string, which is O(T * L) room. The trie shares them.
• Arrays or maps for children? Arrays are fast but use 26 slots per node. Maps save memory when most slots are empty.

Remember this

Grown-up words: the house of hallways is a trie or prefix tree, each room is a node, the doors are child links, the bell is the end of word flag, and shared hallways are shared prefixes.