[LeetCode] 208(LintCode). Implement Trie(Prefix Tree)

Posted 2021-01-27 Push your limit!

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")
>>> false
startsWith("lint")
>>> true
startsWith("linterror")
>>> false
insert("linterror")
search("lintcode)
>>> true
startsWith("linterror")
>>> true

 1 class Trie {
 2     class TrieNode{
 3         private char c;
 4         private Map<Character, TrieNode> map;
 5         private boolean isLeaf;   
 6         
 7         TrieNode() {
 8             this.map = new HashMap<>();
 9         }
10         TrieNode(char c) {
11             this.c = c;
12             this.map = new HashMap<>();
13         }
14     }
15     private TrieNode root;
16     /** Initialize your data structure here. */
17     public Trie() {
18         root = new TrieNode();
19     }
20     
21     /** Inserts a word into the trie. */
22     public void insert(String word) {
23         TrieNode currNode = root;
24         for(int i = 0; i < word.length(); i++) {
25             char currChar = word.charAt(i);
26             if(!currNode.map.containsKey(currChar)) {
27                 TrieNode newNode = new TrieNode(currChar);    
28                 currNode.map.put(currChar, newNode);
29                 currNode = newNode;
30             }
31             else {
32                 currNode = currNode.map.get(currChar);
33             }
34             if(i == word.length() - 1) {
35                 currNode.isLeaf = true;
36             }
37         }
38     }
39     
40     /** Returns if the word is in the trie. */
41     public boolean search(String word) {
42         TrieNode currNode = root;
43         for(int i = 0; i < word.length(); i++) {
44             if(!currNode.map.containsKey(word.charAt(i))) {
45                 return false;
46             }
47             currNode = currNode.map.get(word.charAt(i));
48             if( i == word.length() - 1 && !currNode.isLeaf) {
49                 return false;
50             }
51         }
52         return true;
53     }
54     
55     /** Returns if there is any word in the trie that starts with the given prefix. */
56     public boolean startsWith(String prefix) {
57         TrieNode currNode = root;
58         for(int i = 0; i < prefix.length(); i++) {
59             if(!currNode.map.containsKey(prefix.charAt(i))) {
60                 return false;
61             }
62             currNode = currNode.map.get(prefix.charAt(i));
63         }
64         return true;        
65     }
66 }
67 
68 /**
69  * Your Trie object will be instantiated and called as such:
70  * Trie obj = new Trie();
71  * obj.insert(word);
72  * boolean param_2 = obj.search(word);
73  * boolean param_3 = obj.startsWith(prefix);
74  */