Implement Trie (Prefix Tree)

NeetCode: Implement Trie (Prefix Tree) · LeetCode: 208. Implement Trie (Prefix Tree) · Difficulty: Medium

Implement a Trie with insert(word), search(word) (is exactly this word stored?), and startsWith(prefix) (is any stored word starting with it?). Words use lowercase letters a to z.

Example

Text

insert("dog")
search("dog")      -> true
search("do")       -> false
startsWith("do")   -> true
insert("do")
search("do")       -> true

Brute force: keep a list of words

Store every word in a List<String> (an ArrayList) and scan it for each query. List.contains compares strings with equals.

Java

import java.util.ArrayList;
import java.util.List;

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

    public Trie() {}

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

    public boolean search(String word) {
        return words.contains(word);
    }

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

public class Main {
    public static void main(String[] args) {
        Trie trie = new Trie();
        trie.insert("dog");
        System.out.println(trie.search("dog"));
        System.out.println(trie.search("do"));
        System.out.println(trie.startsWith("do"));
        trie.insert("do");
        System.out.println(trie.search("do"));
    }
}

Output

true
false
true
true

Time: O(N · L) per query for N stored words of length L. Space: O(N · L).

Optimal: a trie

A trie is a tree where each edge is a letter and the path from the root spells a prefix. Each node has an array of 26 child references (null when absent) and a flag marking that a word ends there. The child for letter c is at index c - 'a', because char values are numbers and 'a' to 'z' are consecutive. Insert, search and prefix lookup each walk one node per letter, so the cost depends on the query's length, not on how many words are stored. Words that share a prefix share the nodes for it.

A character outside a to z would give an index outside 0..25 and throw ArrayIndexOutOfBoundsException, so the input must be lowercase letters.

Java

class TrieNode {
    TrieNode[] children = new TrieNode[26];
    boolean isEnd;
}

class Trie {
    private final TrieNode root = new TrieNode();

    public Trie() {}

    public void insert(String word) {
        TrieNode node = root;
        for (int i = 0; i < word.length(); i++) {
            int c = word.charAt(i) - 'a';
            if (node.children[c] == null) node.children[c] = new TrieNode();
            node = node.children[c];
        }
        node.isEnd = true;
    }

    // Returns the node reached by following prefix, or null if the path does not exist.
    private TrieNode find(String prefix) {
        TrieNode node = root;
        for (int i = 0; i < prefix.length(); i++) {
            node = node.children[prefix.charAt(i) - 'a'];
            if (node == null) return null;
        }
        return node;
    }

    public boolean search(String word) {
        TrieNode node = find(word);
        return node != null && node.isEnd;
    }

    public boolean startsWith(String prefix) {
        return find(prefix) != null;
    }
}

public class Main {
    public static void main(String[] args) {
        Trie trie = new Trie();
        trie.insert("dog");
        System.out.println(trie.search("dog"));
        System.out.println(trie.search("do"));
        System.out.println(trie.startsWith("do"));
        trie.insert("do");
        System.out.println(trie.search("do"));
        System.out.println(trie.startsWith("cat"));
    }
}

Output

true
false
true
true
false

Time: O(L) per operation for a word of length L. Space: O(total letters), with up to 26 child slots per node.