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.

implement-trie-leetcode-java

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;
    }
}
Category >> Algorithms >> Interview >> Java  
If you want someone to read your code, please put the code inside <pre><code> and </code></pre> tags. For example:
<pre><code> 
String foo = "bar";
</code></pre>
  • mouniga

    What are the numbers in blue next to each node?

  • Sunny

    I meant it create the tree iteratively not recursively.

  • Sunny

    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

  • Ankit Shah

    In the above diagram how would a word “tennis” fit in, we already have “ten” and n is the leaf node.

  • paul

    simply insert “$” + key in that case..

  • ryanlr

    Right, thanks!

  • Ilya Klyuchnikov

    This code doesn’t handle empty strings

  • DPGraph

    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);
    }

    }