[LintCode] Implement Trie

Problem

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

Example

insert("lintcode")
search("code") // return false
startsWith("lint") // return true
startsWith("linterror") // return false
insert("linterror")
search("lintcode) // return true
startsWith("linterror") // return true

Note

You may assume that all inputs are consist of lowercase letters a-z.

Solution

Insert. To insert a word(String) into a trie, by running every digit of the word in a for-loop, we have to first copy the TrieNode root into a new TrieNode cur like the way we treat TreeNode. Then we have to check if cur.children is null. If so, we create a new children with space of 26 for 26 characters from a to z. Then we check if the i-th children of cur referring to the i-th character of the word is null. If it’s null, put the character in that position, if it’s not, cur goes to the next. When the loop goes to the last digit of the word, set cur.exist to true.

Search. If we inserted several words in a trie, some branches, corresponding to some words, will have the exist tag being true. So, we still have to loop the word. Just like insert operation, if we see cur.children is null or cur.children[index] is null, the word does not exist. However, if the last-children.exist is true, the word exists.

Prefix. To check if a trie contains any words with a certain prefix, it’s the same as search. The only difference is you don’t need to check the extra digits after the prefix.

class TrieNode {
    // Initialize your data structure here.
    boolean exist;
    char ch;
    TrieNode[] children;
    public TrieNode() {   
    }
    public TrieNode(char ch) {
        this.ch = ch;
    }
}

public class Solution {
    private TrieNode root;
    public Solution() {
        root = new TrieNode();
    }
    // Inserts a word into the trie.
    public void insert(String word) {
        if (word == null || word.length() == 0) {
            return;
        }
        TrieNode pre = root;
        for (int i = 0; i < word.length(); i++) {
           if (pre.children == null) {
               pre.children = new TrieNode[26];
           }
           int index = word.charAt(i) - 'a';
           if (pre.children[index] == null) {
               pre.children[index] = new TrieNode(word.charAt(i));
           }
           pre = pre.children[index];
           if (i == word.length() -1) {
               pre.exist = true;
           }
        }
    }
    // Returns if the word is in the trie.
    public boolean search(String word) {
        if (word == null || word.length() == 0) {
            return false;
        }
        TrieNode pre = root;
        for (int i = 0; i < word.length(); i++) {
            int index = word.charAt(i)-'a';
            if (pre.children == null || pre.children[index] == null) {
                return false;
            } 
            if (i == word.length() - 1 && pre.children[index].exist == false) {
                return false;
            }
            pre = pre.children[index];
        }
        return true;
    }

    // Returns if there is any word in the trie
    // that starts with the given prefix.
    public boolean startsWith(String prefix) {
        if (prefix == null || prefix.length() == 0) {
            return false;
        }
        TrieNode pre = root;
        for (int i = 0; i < prefix.length(); i++) {
            int index = prefix.charAt(i)-'a';
            if (pre.children == null || pre.children[index] == null) {
                return false;
            }
            pre = pre.children[index];
        }
        return true;
    }
}