Implement Trie (Prefix Tree) LeetCode solution in C++

1
class Trie {
2
 public:
3
  struct tr {
4
    bool isword;
5
    vector<tr *> next;
6
    tr() {
7
      next.resize(26, NULL); // one pointer for each character
8
      isword = false;
9
    }
10
  };
11

12
  tr *root;
13
  Trie() {
14
    root = new tr();
15
  }
16

17
  void insert(string word) {
18
    tr *cur = root;
19

20
    for (auto c : word) {
21
      // if pointer to char doesn't exist create one
22
      if (cur->next[c - 'a'] == NULL) {
23
        cur->next[c - 'a'] = new tr();
24
      }
25
      cur = cur->next[c - 'a'];
26
    }
27

28
    // set as end or valid word
29
    cur->isword = true;
30
  }
31

32
  bool search(string word) {
33
    tr *cur = root;
34

35
    for (auto c : word) {
36
      // if pointer to char doesn't exist then word is not present
37
      if (cur->next[c - 'a'] == NULL) {
38
        return false;
39
      }
40

41
      cur = cur->next[c - 'a'];
42
    }
43

44
    // check if it is prefix or a word
45
    if (cur->isword) {
46
      return true;
47
    }
48
    return false;
49
  }
50

51
  // similar to search
52
  // here we don't check if it is end or not
53
  bool startsWith(string prefix) {
54
    tr *cur = root;
55

56
    for (auto c : prefix) {
57
      if (cur->next[c - 'a'] == NULL) {
58
        return false;
59
      }
60

61
      cur = cur->next[c - 'a'];
62
    }
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
 * bool param_2 = obj->search(word);
73
 * bool param_3 = obj->startsWith(prefix);
74
 */