Implement Trie (Prefix Tree)

LeetCode Q 208 - Implement Trie (Prefix Tree)

Implement a trie with insert, search, and startsWith methods.

Example:
Trie trie = new Trie();
trie.insert(“apple”);
trie.search(“apple”); // returns true
trie.search(“app”); // returns false
trie.startsWith(“app”); // returns true
trie.insert(“app”);
trie.search(“app”); // returns true

Note:

• You may assume that all inputs are consist of lowercase letters a-z.
• All inputs are guaranteed to be non-empty strings.

Solution

Code:

class TrieNode {
	TrieNode[] children;
	boolean isEnd;
	public TrieNode () { this.children = new TrieNode[26]; }
}

// Initialize your data structure here. 
TrieNode root;
public Trie() {
	root = new TrieNode();
}

// Inserts a word into the trie.
public void insert(String word) {
	
	if (word == null || word.length() == 0) return;
	
	TrieNode node = root;
	
	for (char ch: word.toCharArray()) {
		if (node.children[ch - 'a'] == null)
			node.children[ch - 'a'] = new TrieNode();
		node = node.children[ch - 'a'];
	}
	
	node.isEnd = true;
	
	return;
}

// Returns if the word is in the trie. 
public boolean search(String word) {
	
	TrieNode node = root;
	
	for (char ch: word.toCharArray()) {
		if (node.children[ch - 'a'] == null)
			return false;
		node = node.children[ch - 'a'];
	}
	
	return node.isEnd;
}

// Returns if there is any word in the trie that starts with the given prefix.
public boolean startsWith(String prefix) {

	TrieNode node = root;
	
	for (char ch: prefix.toCharArray()) {
		if (node.children[ch - 'a'] == null)
			return false;
		node = node.children[ch - 'a'];
	}
	
	return true;
}