Count Pairs With XOR in a Range LeetCode solution in Java

1
class Solution {
2
  public int countPairs(int[] nums, int low, int high) {
3
    Trie trie = new Trie();
4
    int cnt = 0;
5
    for (int i = nums.length - 1; i >= 0; i--) {
6
      // count all the element whose xor is less the low
7
      int cnt1 = trie.maxXor(nums[i], low);
8
      // count all the element whose xor is less the high+1
9
      int cnt2 = trie.maxXor(nums[i], high + 1);
10
      trie.add(nums[i]);
11
      cnt += cnt2 - cnt1;
12
    }
13
    return cnt;
14
  }
15
}
16

17
class Trie {
18
  private Node root;
19

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

24
  public void add(int x) {
25
    Node cur = root;
26
    for (int i = 31; i >= 0; i--) {
27
      int bit = (x >> i) & 1;
28
      if (!cur.contains(bit)) {
29
        cur.put(bit);
30
      }
31
      cur.inc(bit);
32
      cur = cur.get(bit);
33
    }
34
  }
35

36
  public int maxXor(int x, int limit) {
37
    int low_cnt = 0;
38
    Node cur = root;
39
    for (int i = 31; i >= 0 && cur != null; i--) {
40
      int bit = (x >> i) & (1);
41
      int req = (limit >> i) & 1;
42
      if (req == 1) {
43
        if (cur.contains(bit)) {
44
          low_cnt += cur.get(bit).cnt;
45
        }
46
        cur = cur.get(1 - bit);
47
      } else {
48
        cur = cur.get(bit);
49
      }
50
    }
51
    return low_cnt;
52
  }
53
}
54

55
class Node {
56
  private Node links[];
57
  int cnt;
58

59
  Node() {
60
    links = new Node[2];
61
    cnt = 0;
62
  }
63

64
  public void put(int bit) {
65
    links[bit] = new Node();
66
  }
67

68
  public Node get(int bit) {
69
    return links[bit];
70
  }
71

72
  public boolean contains(int bit) {
73
    return links[bit] != null;
74
  }
75

76
  public void inc(int bit) {
77
    links[bit].cnt++;
78
  }
79
}