LeetCode – Implement Trie (Prefix Tree) (Java)
Implement a trie with insert, search, and startsWith methods.
Java Solution 1
A trie node should contains the character, its children and the flag that marks if it is a leaf node. You can use this diagram to walk though the Java solution.
class TrieNode { char c; HashMap<Character, TrieNode> children = new HashMap<Character, TrieNode>(); boolean isLeaf; public TrieNode() {} public TrieNode(char c){ this.c = c; } } |
public class Trie { private TrieNode root; public Trie() { root = new TrieNode(); } // Inserts a word into the trie. public void insert(String word) { HashMap<Character, TrieNode> children = root.children; for(int i=0; i<word.length(); i++){ char c = word.charAt(i); TrieNode t; if(children.containsKey(c)){ t = children.get(c); }else{ t = new TrieNode(c); children.put(c, t); } children = t.children; //set leaf node if(i==word.length()-1) t.isLeaf = true; } } // Returns if the word is in the trie. public boolean search(String word) { TrieNode t = searchNode(word); if(t != null && t.isLeaf) return true; else return false; } // Returns if there is any word in the trie // that starts with the given prefix. public boolean startsWith(String prefix) { if(searchNode(prefix) == null) return false; else return true; } public TrieNode searchNode(String str){ Map<Character, TrieNode> children = root.children; TrieNode t = null; for(int i=0; i<str.length(); i++){ char c = str.charAt(i); if(children.containsKey(c)){ t = children.get(c); children = t.children; }else{ return null; } } return t; } } |
Java Solution 2 - Improve Performance by Using an Array
Each trie node can only contains 'a'-'z' characters. So we can use a small array to store the character.
class TrieNode { TrieNode[] arr; boolean isEnd; // Initialize your data structure here. public TrieNode() { this.arr = new TrieNode[26]; } } public class Trie { private TrieNode root; public Trie() { root = new TrieNode(); } // Inserts a word into the trie. public void insert(String word) { TrieNode p = root; for(int i=0; i<word.length(); i++){ char c = word.charAt(i); int index = c-'a'; if(p.arr[index]==null){ TrieNode temp = new TrieNode(); p.arr[index]=temp; p = temp; }else{ p=p.arr[index]; } } p.isEnd=true; } // Returns if the word is in the trie. public boolean search(String word) { TrieNode p = searchNode(word); if(p==null){ return false; }else{ if(p.isEnd) return true; } return false; } // Returns if there is any word in the trie // that starts with the given prefix. public boolean startsWith(String prefix) { TrieNode p = searchNode(prefix); if(p==null){ return false; }else{ return true; } } public TrieNode searchNode(String s){ TrieNode p = root; for(int i=0; i<s.length(); i++){ char c= s.charAt(i); int index = c-'a'; if(p.arr[index]!=null){ p = p.arr[index]; }else{ return null; } } if(p==root) return null; return p; } } |
<pre><code> String foo = "bar"; </code></pre>
Leave a comment
A recursive Java version:
class TrieNode {
Object val;
HashMap map;
// Initialize your data structure here.
public TrieNode() {
val = null;
map = new HashMap();
}
}
public class Trie {
private TrieNode root;
public Trie() {
root = new TrieNode();
}
// Inserts a word into the trie.
public void insert(String word) {
insert(word, root, 0);
}
private TrieNode insert(String word, TrieNode x, int pos) {
if(x == null) x = new TrieNode();
if(pos == word.length()) {
x.val = new Object();
return x;
}
char c = word.charAt(pos);
x.map.put(c, insert(word, x.map.get(c), pos+1));
return x;
}
// Returns if the word is in the trie.
public boolean search(String word) {
return search(word, root, 0);
}
private boolean search(String word, TrieNode x, int pos) {
if(x == null) return false;
if(pos == word.length()) {
return x.val != null;
}
return search(word, x.map.get(word.charAt(pos)), pos+1);
}
// Returns if there is any word in the trie
// that starts with the given prefix.
public boolean startsWith(String prefix) {
return startWith(prefix, root, 0);
}
private boolean startWith(String prefix, TrieNode x, int pos) {
if(x == null) return false;
if(pos == prefix.length()) return x != null;
return startWith(prefix, x.map.get(prefix.charAt(pos)), pos+1);
}
}
This code doesn’t handle empty strings
Right, thanks!
simply insert “$” + key in that case..
In the above diagram how would a word “tennis” fit in, we already have “ten” and n is the leaf node.
The insert method just adds the characters in the word into the HashMap. It does not create a tree. Look at the alternative implementation here https://github.com/brianfromoregon/trie/blob/master/src/main/java/net/bcharris/trie/TrieImpl.java#L154
I meant it create the tree iteratively not recursively.
What are the numbers in blue next to each node?
In this public TrieNode searchNode(String str)
Its a has map
It marks the leaf, i.e.- Word ends here from root .