leetcode-study/implement-trie-prefix-tree/sangbeenmoon.java at 4e7811fb97304781070100b5dda9b594121ffe71 · DaleStudy/leetcode-study · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
import java.util.HashMap;
import java.util.Map;

class Trie {

    Map<Character, Trie> map;
    boolean isEnd = false;

    public Trie() {
        map = new HashMap<>();
    }

    public void insert(String word) {
        Trie trie;
        if (map.containsKey(word.charAt(0))) {
            trie = map.get(word.charAt(0));
        } else {
            trie = new Trie();
        }
        map.put(word.charAt(0), trie);

        if (word.length() == 1) {
            map.get(word.charAt(0)).isEnd = true;
            return;
        }

        trie.insert(word.substring(1));
    }

    public boolean search(String word) {
        if (map.containsKey(word.charAt(0))) {
            if (word.length() == 1) {
                return map.get(word.charAt(0)).isEnd;
            }
            return map.get(word.charAt(0)).search(word.substring(1));
        }
        return false;
    }

    public boolean startsWith(String prefix) {
        if (map.containsKey(prefix.charAt(0))) {
            if (prefix.length() == 1) {
                return true;
            }
            return map.get(prefix.charAt(0)).startsWith(prefix.substring(1));
        }
        return false;
    }
}

/**
 * Your Trie object will be instantiated and called as such:
 * Trie obj = new Trie();
 * obj.insert(word);
 * boolean param_2 = obj.search(word);
 * boolean param_3 = obj.startsWith(prefix);
 */