Count Triplets That Can Form Two Arrays of Equal XOR LeetCode solution in C++

1
class TrieNode {
2
 public:
3
  int numOfIndex;
4
  int sumOfIndex;
5
  TrieNode *child[2];
6

7
  TrieNode() : numOfIndex(0), sumOfIndex(0) {
8
    child[0] = NULL;
9
    child[1] = NULL;
10
  }
11
};
12
class Solution {
13
 public:
14
  void addNumber(TrieNode *root, int num, int idx) {
15
    for (int i = 31; i >= 0; i--) {
16
      int bit = 1 & (num >> i);
17

18
      if (root->child[bit] == NULL) {
19
        root->child[bit] = new TrieNode();
20
      }
21
      root = root->child[bit];
22
    }
23
    root->sumOfIndex += idx;
24
    root->numOfIndex++;
25
  }
26
  int calculateIndexPair(TrieNode *root, int num, int idx) {
27
    for (int i = 31; i >= 0; i--) {
28
      int bit = 1 & (num >> i);
29

30
      if (root->child[bit] == NULL) {
31
        return 0;
32
      }
33
      root = root->child[bit];
34
    }
35
    return (((root->numOfIndex) * idx) - (root->sumOfIndex));
36
  }
37
  int countTriplets(vector<int> &arr) {
38
    long long ans = 0;
39
    int XOR = 0;
40
    TrieNode *root = new TrieNode();
41

42
    for (int i = 0; i < arr.size(); i++) {
43
      addNumber(root, XOR, i);
44
      XOR ^= arr[i];
45
      ans = (ans + calculateIndexPair(root, XOR, i)) % 1000000007;
46
    }
47
    return ans;
48
  }
49
};