LeetCode: Implement Trie (Prefix Tree)

LeetCode: Implement Trie (Prefix Tree)

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

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
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
class TrieNode {

private boolean isLeaf;
private Map<Character, TrieNode> children;
private char value;

// Initialize your data structure here.
public TrieNode() {
this.isLeaf = false;
this.children = new HashMap<Character, TrieNode>();
}

public TrieNode(char value) {
this();
this.value = value;
}

public void putChild(char c) {
this.children.put(c, new TrieNode(c));
}
public boolean containsChild(char c) {
return this.children.containsKey(c);
}
public TrieNode getChild(char c) {
return containsChild(c)? this.children.get(c) : null;
}

public void setLeaf() {
this.isLeaf = true;
}

public boolean isLeaf() {
return this.isLeaf;
}
}

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.isEmpty()) {
return;
}

char[] chars = word.toCharArray();
TrieNode tmp = root;
for(char c : chars) {
if(!tmp.containsChild(c)) {
tmp.putChild(c);
}
tmp = tmp.getChild(c);
}
tmp.setLeaf();
}

// Returns if the word is in the trie.
public boolean search(String word) {
if(word == null || word.isEmpty()) {
return true;
}
char[] chars = word.toCharArray();
TrieNode tmp = root;
for(char c : chars) {
if(!tmp.containsChild(c)) {
return false;
}
tmp = tmp.getChild(c);
}
return tmp.isLeaf();
}

// Returns if there is any word in the trie
// that starts with the given prefix.
public boolean startsWith(String prefix) {
if(prefix == null || prefix.isEmpty()) {
return true;
}
char[] chars = prefix.toCharArray();
TrieNode tmp = root;
for(char c : chars) {
if(!tmp.containsChild(c)) {
return false;
}
tmp = tmp.getChild(c);
}
return true;
}
}