Maximum Genetic Difference Query LeetCode solution in C++

1
/*
2
    Time: O(q+n)
3
    Space: O(q+n)
4
    Tag: Trie, N-ary Tree
5
    Difficulty: H
6
*/
7
struct Node {
8
  int count;
9
  vector<Node *> children;
10
  Node() {
11
    children.resize(2, NULL);
12
    count = 0;
13
  }
14
};
15

16
class Trie {
17
  Node *root;
18

19
 public:
20
  Trie() {
21
    root = new Node();
22
  }
23

24
  void insert(int num) {
25
    Node *cur = root;
26
    for (int i = 18; i >= 0; i--) {
27
      bool bit = (num & (1 << i));
28
      if (!cur->children[bit]) cur->children[bit] = new Node();
29
      cur = cur->children[bit];
30
      cur->count++;
31
    }
32
  }
33

34
  void remove(int num) {
35
    Node *cur = root;
36
    for (int i = 18; i >= 0; i--) {
37
      bool bit = (num & (1 << i));
38
      cur = cur->children[bit];
39
      cur->count--;
40
    }
41
  }
42

43
  int search(int num) {
44
    int res = 0;
45
    Node *cur = root;
46
    for (int i = 18; i >= 0; i--) {
47
      bool bit = (num & (1 << i));
48
      if (cur->children[!bit] && cur->children[!bit]->count > 0) {
49
        res += (1 << i);
50
        cur = cur->children[!bit];
51
      } else {
52
        cur = cur->children[bit];
53
      }
54
    }
55
    return res;
56
  }
57
};
58

59
class Solution {
60
  vector<vector<pair<int, int>>> mp;
61
  Trie *t;
62
  vector<int> res;
63

64
  void dfs(vector<int> gr[], int cur) {
65
    t->insert(cur);
66
    for (auto &child : mp[cur]) {
67
      res[child.second] = t->search(child.first);
68
    }
69
    for (int &child : gr[cur]) {
70
      dfs(gr, child);
71
    }
72
    t->remove(cur);
73
  }
74

75
 public:
76
  vector<int> maxGeneticDifference(vector<int> &parents, vector<vector<int>> &queries) {
77
    t = new Trie();
78

79
    int n = parents.size();
80
    vector<int> gr[n];
81
    res.resize(queries.size());
82

83
    int root = -1;
84

85
    for (int i = 0; i < n; i++) {
86
      if (parents[i] == -1)
87
        root = i;
88
      else
89
        gr[parents[i]].push_back(i);
90
    }
91
    mp.resize(n);
92
    for (int i = 0; i < queries.size(); i++) {
93
      mp[queries[i][0]].push_back({queries[i][1], i});
94
    }
95

96
    dfs(gr, root);
97

98
    return res;
99
  }
100
};