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 the trie in the following 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;
    }
}

If the same words can be inserted more than once, what do you need to change to make it work?

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>

Kishore Srinivas

to find whether there is a next word (link) to the trieNode.
Anshika

What is the use of keeping char c in TrieNode class?
Hitesh Gupta

It marks the leaf, i.e.- Word ends here from root .
Tarun

In this public TrieNode searchNode(String str)

Its a has map
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); } }

LeetCode – Implement Trie (Prefix Tree) (Java)

Related posts: