Implement Trie (Prefix Tree) 解答

Question

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

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

Solution

A trie node should contains the character, its children and the flag that marks if it is a leaf node.

Trie is an efficient information retrieval data structure. Using trie, search complexities can be brought to optimal limit (key length).

Every node of trie consists of multiple branches. Each branch represents a possible character of keys. We need to mark the last node of every key as leaf node.

We design TrieNode to store 1. value 2. isLeaf 3. a map of children (map is used for quick selection)

In this way, we can implement Trie start from root. Add new child to original parent is implemented by put action in map.

Note here, search and startsWith differs.

　　null

　　 |

　　 a

　　 |

　　 b

for "a", search will return false, while startsWith will return true.

/**
 * 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 {
    char val;
    boolean isLast;
    Map<Character, TrieNode> children;
    // Initialize your data structure here.
    public TrieNode(char val) {
        children = new HashMap<Character, TrieNode>();
        this.val = val;
        this.isLast = false;
    }
}

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 || word.length() == 0) {
            return;
        }
        TrieNode p = this.root;
        char[] arr = word.toCharArray();
        for (char cur : arr) {
            if (!p.children.containsKey(cur)) {
                p.children.put(cur, new TrieNode(cur));
            }
            p = p.children.get(cur);
        }
        p.isLast = true;
    }

    // Returns if the word is in the trie.
    public boolean search(String word) {
        if (word == null || word.length() == 0) {
            return false;
        }
        char[] arr = word.toCharArray();
        TrieNode p = this.root;
        for (char cur : arr) {
            Map<Character, TrieNode> children = p.children;
            if (!children.containsKey(cur)) {
                return false;
            }
            p = children.get(cur);
        }
        return p.isLast;
    }

    // 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;
        }
        char[] arr = prefix.toCharArray();
        TrieNode p = this.root;
        for (char cur : arr) {
            Map<Character, TrieNode> children = p.children;
            if (!children.containsKey(cur)) {
                return false;
            }
            p = children.get(cur);
        }
        return true;
    }
}