Implement Trie 442

Question

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

Notice

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

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

Solution

TrieNode class：中用一个HashMap来记录紧跟该node之后的所有node（每个字符都是一个node，即该字符之后所有可能字符）。若该node是叶节点，则将从root到该node路径上所有字符组成的字符串保存在该node的s中，并标记该路径是一个string。

insert（）：将字符串插入Trie中。依次取出每一位字符，从root开始，若当前node之后的node不包含该字符，则新建一个该字符的node并加入当前node的suntree中，然后将当前node移往新建的node，循环直到字符串尾。

search（）：搜索Trie中是否含有该字符串。依次取出每一位字符，从root开始，若当前node之后的node不包含该字符，则返回false，直到字符串尾。此时，若当前节点为叶节点，则返回true，否则返回false。

startWith（）：搜索Trie中是否含有以prefix为开头的字符串。方法和search（）一样，最后不用判断是否当前节点为叶节点。

代码如下：

/**
 * Your Trie object will be instantiated and called as such:
 * Trie trie = new Trie();
 * trie.insert("lintcode");
 * trie.search("lint"); will return false
 * trie.startsWith("lint"); will return true
 */
class TrieNode {
    String s;
    boolean isString;
    HashMap<Character, TrieNode> subTree;
    // Initialize your data structure here.
    public TrieNode() {
        s = "";
        isString = false;
        subTree = new HashMap<Character, TrieNode>();
    }
}

public class Trie {
    private TrieNode root;

    public Trie() {
        root = new TrieNode();
    }

    // Inserts a word into the trie.
    public void insert(String word) {
        if(word == null){
            return;
        }

        TrieNode now = root;
        for(int i = 0; i < word.length(); i++){
            if(!now.subTree.containsKey(word.charAt(i))){
                now.subTree.put(word.charAt(i), new TrieNode());
            }
            now = now.subTree.get(word.charAt(i));
        }
        now.s = word;
        now.isString = true;
    }

    // Returns if the word is in the trie.
    public boolean search(String word) {
        if(word == null){
            return true;
        }

        TrieNode now = root;
        for(int i = 0; i < word.length(); i++){
            if(!now.subTree.containsKey(word.charAt(i))){
                return false;
            }
            now = now.subTree.get(word.charAt(i));
        }

        if(now.isString){
            return true;
        }
        return false;
    }

    // Returns if there is any word in the trie
    // that starts with the given prefix.
    public boolean startsWith(String prefix) {
        if(prefix == null){
            return true;
        }

        TrieNode now = root;
        for(int i = 0; i < prefix.length(); i++){
            if(!now.subTree.containsKey(prefix.charAt(i))){
                return false;
            }
            now = now.subTree.get(prefix.charAt(i));
        }
        return true;
    }
}