Implement Trie (Prefix Tree)

Medium Problem 95 of 151 Java and Python

Problem

A trie (prefix tree) stores strings so that lookups by prefix are fast. Implement the Trie class:

insert(word) adds a word, search(word) returns true if the exact word was inserted, and startsWith(prefix) returns true if any inserted word starts with prefix.

Example 1

Input: insert("apple"), search("apple"), search("app"), startsWith("app"), insert("app"), search("app")
Output: true, false, true, true

Constraints

1 ≤ word.length, prefix.length ≤ 2000
Only lowercase English letters.
At most 3 × 10⁴ calls in total.

Explanation

Words that share a prefix share a path from the root. Each node represents one prefix and has up to 26 children, one per next letter. A flag on a node says "a word ends here".

All three operations just walk down the path letter by letter, so they cost O(length of the word), regardless of how many words are stored.

Approach 1 · Naive

Set of words

Time O(L) insert and search, O(N · L) startsWith
Space O(total characters)

Store words in a hash set. search is O(L), but startsWith must scan every stored word.

insert: add to the set.
search: set lookup.
startsWith: check every word for the prefix.

Java

import java.util.*;

class Trie {
    private final Set<String> words = new HashSet<>();

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

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

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

Python

class Trie:
    def __init__(self):
        self.words = set()

    def insert(self, word: str) -> None:
        self.words.add(word)

    def search(self, word: str) -> bool:
        return word in self.words

    def startsWith(self, prefix: str) -> bool:
        return any(word.startswith(prefix) for word in self.words)

Approach 2 · Optimal

Trie nodes with 26 children

Time O(L) per operation
Space O(total characters × 26)

Each node has an array of 26 child pointers and an end flag.

insert: walk from the root, creating missing children; mark the last node end.
search: walk; fail on a missing child; succeed only if the last node is end.
startsWith: walk; succeed if the whole prefix exists.

Java

class Trie {
    private static class Node {
        Node[] children = new Node[26];
        boolean end;
    }

    private final Node root = new Node();

    public void insert(String word) {
        Node node = root;
        for (char c : word.toCharArray()) {
            if (node.children[c - 'a'] == null) {
                node.children[c - 'a'] = new Node();
            }
            node = node.children[c - 'a'];
        }
        node.end = true;
    }

    public boolean search(String word) {
        Node node = walk(word);
        return node != null && node.end;
    }

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

    private Node walk(String s) {
        Node node = root;
        for (char c : s.toCharArray()) {
            node = node.children[c - 'a'];
            if (node == null) {
                return null;
            }
        }
        return node;
    }
}

Python

class TrieNode:
    def __init__(self):
        self.children = {}
        self.end = False


class Trie:
    def __init__(self):
        self.root = TrieNode()

    def insert(self, word: str) -> None:
        node = self.root
        for c in word:
            node = node.children.setdefault(c, TrieNode())
        node.end = True

    def _walk(self, s):
        node = self.root
        for c in s:
            node = node.children.get(c)
            if node is None:
                return None
        return node

    def search(self, word: str) -> bool:
        node = self._walk(word)
        return node is not None and node.end

    def startsWith(self, prefix: str) -> bool:
        return self._walk(prefix) is not None

Complexity compared

Approach	Time	Space
Naive Set of words	O(L) insert and search, O(N · L) startsWith	O(total characters)
Optimal Trie nodes with 26 children	O(L) per operation	O(total characters × 26)

Key takeaways

search and startsWith differ only in whether the final node must be marked as the end of a word.
A dictionary of children (Python) uses less memory than a fixed array of 26 when the alphabet is sparse.

Your notes

Saved in this browser as you type.

Solved it on your own?

Mark it done to track your progress, or star it to come back before an interview.

Next problem