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?

16 thoughts on “LeetCode – Implement Trie (Prefix Tree) (Java)”

  1. You are right,
    From implementation perspective ,the char is not really needed (it is just use as a mark of which character this node represent).
    Yet, it is great post.

  2. I disagree that `char c;` is used for anything:
    1) in the hand drawn diagram at the top of the article the root node stores no character. So ‘c’ is not needed.
    2) you’re using children.containsKey(c) for insert and search methods. So ‘c’ is not needed.

    However, `char c;` is useful to detect when someone blindly copy/paste’s the code

  3. In insert in second implementation here is no need of else statement.Can be optimized like given below.

    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=p.arr[index];

    }

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

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

    }

Leave a Comment