Card 04/ 07
Trie in Python, Java and TypeScript
Insert, search and starts-with, with children kept in a map at each node.
None of Python, Java or TypeScript ships a trie in its standard library, so a trie is a short class written by hand. This card covers that class in all three languages: a node type, insert, search, a prefix check, and a method that lists every word under a prefix, with each node's children kept in a map from letter to child.
A trie is a tree whose root holds no letter and whose every other node stands for one letter; a node carries an end-of-word mark when its path from the root spells a stored word. A map, also called a dictionary, stores values under keys and finds a key in about constant time on average.
All three versions share one shape. A private find method walks the letters of a string and returns the node it reaches, or nothing if a letter has no child. Search and the prefix check both call it: search also requires the mark, the prefix check does not. Listing words finds the prefix's node, then visits every node below it, collecting each marked path. Each listing inserts sun, sum, sung and sky and makes the same six calls.
The same trie in each language
class Node:
def __init__(self):
self.children = {}
self.is_word = False
class Trie:
def __init__(self):
self.root = Node()
def insert(self, word):
node = self.root
for letter in word:
node = node.children.setdefault(letter, Node())
node.is_word = True
def search(self, word):
node = self._find(word)
return node is not None and node.is_word
def starts_with(self, prefix):
return self._find(prefix) is not None
def words_with_prefix(self, prefix):
node = self._find(prefix)
words = []
if node is not None:
self._collect(node, prefix, words)
return words
def _find(self, text):
node = self.root
for letter in text:
node = node.children.get(letter)
if node is None:
return None
return node
def _collect(self, node, text, words):
if node.is_word:
words.append(text)
for letter in sorted(node.children):
self._collect(node.children[letter], text + letter, words)
trie = Trie()
for word in ["sun", "sum", "sung", "sky"]:
trie.insert(word)
print("search sun:", trie.search("sun"))
print("search su:", trie.search("su"))
print("starts_with su:", trie.starts_with("su"))
print("starts_with so:", trie.starts_with("so"))
print("words with prefix su:", trie.words_with_prefix("su"))
print("words with prefix s:", trie.words_with_prefix("s"))search sun: True
search su: False
starts_with su: True
starts_with so: False
words with prefix su: ['sum', 'sun', 'sung']
words with prefix s: ['sky', 'sum', 'sun', 'sung']import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;
public class TrieCode {
static final class Node {
final Map<Character, Node> children = new TreeMap<>();
boolean isWord;
}
static final class Trie {
private final Node root = new Node();
void insert(String word) {
Node node = root;
for (char letter : word.toCharArray()) {
node = node.children.computeIfAbsent(letter, key -> new Node());
}
node.isWord = true;
}
boolean search(String word) {
Node node = find(word);
return node != null && node.isWord;
}
boolean startsWith(String prefix) {
return find(prefix) != null;
}
List<String> wordsWithPrefix(String prefix) {
List<String> words = new ArrayList<>();
Node node = find(prefix);
if (node != null) {
collect(node, new StringBuilder(prefix), words);
}
return words;
}
private Node find(String text) {
Node node = root;
for (char letter : text.toCharArray()) {
node = node.children.get(letter);
if (node == null) {
return null;
}
}
return node;
}
private void collect(Node node, StringBuilder text, List<String> words) {
if (node.isWord) {
words.add(text.toString());
}
for (Map.Entry<Character, Node> child : node.children.entrySet()) {
text.append(child.getKey());
collect(child.getValue(), text, words);
text.deleteCharAt(text.length() - 1);
}
}
}
public static void main(String[] args) {
Trie trie = new Trie();
for (String word : List.of("sun", "sum", "sung", "sky")) {
trie.insert(word);
}
System.out.println("search sun: " + trie.search("sun"));
System.out.println("search su: " + trie.search("su"));
System.out.println("startsWith su: " + trie.startsWith("su"));
System.out.println("startsWith so: " + trie.startsWith("so"));
System.out.println("words with prefix su: " + trie.wordsWithPrefix("su"));
System.out.println("words with prefix s: " + trie.wordsWithPrefix("s"));
}
}search sun: true
search su: false
startsWith su: true
startsWith so: false
words with prefix su: [sum, sun, sung]
words with prefix s: [sky, sum, sun, sung]class TrieNode {
readonly children = new Map<string, TrieNode>();
isWord = false;
}
class Trie {
private readonly root = new TrieNode();
insert(word: string): void {
let node = this.root;
for (const letter of word) {
let child = node.children.get(letter);
if (child === undefined) {
child = new TrieNode();
node.children.set(letter, child);
}
node = child;
}
node.isWord = true;
}
search(word: string): boolean {
const node = this.find(word);
return node !== undefined && node.isWord;
}
startsWith(prefix: string): boolean {
return this.find(prefix) !== undefined;
}
wordsWithPrefix(prefix: string): string[] {
const words: string[] = [];
const node = this.find(prefix);
if (node !== undefined) {
this.collect(node, prefix, words);
}
return words;
}
private find(text: string): TrieNode | undefined {
let node: TrieNode | undefined = this.root;
for (const letter of text) {
node = node.children.get(letter);
if (node === undefined) {
return undefined;
}
}
return node;
}
private collect(node: TrieNode, text: string, words: string[]): void {
if (node.isWord) {
words.push(text);
}
for (const letter of [...node.children.keys()].sort()) {
this.collect(node.children.get(letter) as TrieNode, text + letter, words);
}
}
}
const trie = new Trie();
for (const word of ["sun", "sum", "sung", "sky"]) {
trie.insert(word);
}
console.log(`search sun: ${trie.search("sun")}`);
console.log(`search su: ${trie.search("su")}`);
console.log(`startsWith su: ${trie.startsWith("su")}`);
console.log(`startsWith so: ${trie.startsWith("so")}`);
console.log(`words with prefix su: [${trie.wordsWithPrefix("su").join(", ")}]`);
console.log(`words with prefix s: [${trie.wordsWithPrefix("s").join(", ")}]`);search sun: true
search su: false
startsWith su: true
startsWith so: false
words with prefix su: [sum, sun, sung]
words with prefix s: [sky, sum, sun, sung]The highlighted lines are the ones that differ between a trie and any other tree. Insert creates a missing child and moves to it in one call: setdefault in Python, computeIfAbsent in Java, and a get followed by a set in TypeScript, whose Map has neither. The line after the loop sets the mark. The search returns the mark of the node it reached, while the prefix check returns only whether a node was reached.
The word lists come out in alphabetical order because the children are visited in letter order: Python sorts the keys at each node, Java keeps them in a TreeMap, which iterates its keys in sorted order, and TypeScript sorts them as Python does. A plain dictionary or HashMap would return the same words in an order that depends on insertion or hashing.
Many versions replace the map with an array of 26 slots, one per lowercase letter, indexed by letter - 'a'. That suits a fixed lowercase alphabet: finding a child is an index, and the children come out in order with no sorting. It costs 26 slots in every node, most of them empty, and it fails on any character outside the 26. The map holds only the children that exist and accepts any character, at the cost of hashing each letter. Either way, each operation takes O(L) steps for a string of L letters.
← → move through sections and cards