1707. Maximum XOR With an Element From Array

• Time: $O(\texttt{sort}(n) + \texttt{sort}(q) + q\log_2(\max(\texttt{nums})))$
• Space: $O(n + q)$

C++JavaPython

struct TrieNode {
  vector<shared_ptr<TrieNode>> children;
  TrieNode() : children(2) {}
};

class BitTrie {
 public:
  BitTrie(int maxBit) : maxBit(maxBit) {}

  void insert(int num) {
    shared_ptr<TrieNode> node = root;
    for (int i = maxBit; i >= 0; --i) {
      const int bit = num >> i & 1;
      if (node->children[bit] == nullptr)
        node->children[bit] = make_shared<TrieNode>();
      node = node->children[bit];
    }
  }

  int getMaxXor(int num) {
    int maxXor = 0;
    shared_ptr<TrieNode> node = root;
    for (int i = maxBit; i >= 0; --i) {
      const int bit = num >> i & 1;
      const int toggleBit = bit ^ 1;
      if (node->children[toggleBit] != nullptr) {
        maxXor = maxXor | 1 << i;
        node = node->children[toggleBit];
      } else if (node->children[bit] != nullptr) {
        node = node->children[bit];
      } else {  // Nothing in the Bit Trie.
        return 0;
      }
    }
    return maxXor;
  }

 private:
  const int maxBit;
  shared_ptr<TrieNode> root = make_shared<TrieNode>();
};

class Solution {
 public:
  vector<int> maximizeXor(vector<int>& nums, vector<vector<int>>& queries) {
    vector<int> ans(queries.size(), -1);
    vector<pair<vector<int>, int>> queryAndIndexes;

    const int maxNumInNums = *max_element(nums.begin(), nums.end());
    const int maxNumInQuery =
        max_element(queries.begin(), queries.end(),
                    [](const vector<int>& a, const vector<int>& b) {
      return a[0] < b[0];
        })->at(0);
    const int maxBit = static_cast<int>(log2(max(maxNumInNums, maxNumInQuery)));
    BitTrie bitTrie(maxBit);

    for (int i = 0; i < queries.size(); i++)
      queryAndIndexes.emplace_back(queries[i], i);

    sort(nums.begin(), nums.end());
    sort(queryAndIndexes.begin(), queryAndIndexes.end(),
         [](const pair<vector<int>, int>& a, const pair<vector<int>, int>& b) {
      return a.first[1] < b.first[1];
    });

    int i = 0;  // nums' index
    for (const auto& [query, index] : queryAndIndexes) {
      const int x = query[0];
      const int m = query[1];
      while (i < nums.size() && nums[i] <= m)
        bitTrie.insert(nums[i++]);
      if (i > 0 && nums[i - 1] <= m)
        ans[index] = bitTrie.getMaxXor(x);
    }

    return ans;
  }
};

class TrieNode {
  public TrieNode[] children = new TrieNode[2];
}

class BitTrie {
  public BitTrie(int maxBit) {
    this.maxBit = maxBit;
  }

  public void insert(int num) {
    TrieNode node = root;
    for (int i = maxBit; i >= 0; --i) {
      final int bit = (int) (num >> i & 1);
      if (node.children[bit] == null)
        node.children[bit] = new TrieNode();
      node = node.children[bit];
    }
  }

  public int getMaxXor(int num) {
    int maxXor = 0;
    TrieNode node = root;
    for (int i = maxBit; i >= 0; --i) {
      final int bit = (int) (num >> i & 1);
      final int toggleBit = bit ^ 1;
      if (node.children[toggleBit] != null) {
        maxXor = maxXor | 1 << i;
        node = node.children[toggleBit];
      } else if (node.children[bit] != null) {
        node = node.children[bit];
      } else { // Nothing in the Bit Trie.
        return 0;
      }
    }
    return maxXor;
  }

  private int maxBit;
  private TrieNode root = new TrieNode();
}

class Solution {
  public int[] maximizeXor(int[] nums, int[][] queries) {
    int[] ans = new int[queries.length];
    Arrays.fill(ans, -1);
    Pair<int[], Integer>[] queryAndIndexes = new Pair[queries.length];

    final int maxNumInNums = Arrays.stream(nums).max().getAsInt();
    final int maxNumInQuery = Arrays.stream(queries).mapToInt(query -> query[0]).max().getAsInt();
    final int maxBit = (int) (Math.log(Math.max(maxNumInNums, maxNumInQuery)) / Math.log(2));
    BitTrie bitTrie = new BitTrie(maxBit);

    for (int i = 0; i < queries.length; ++i)
      queryAndIndexes[i] = new Pair<>(queries[i], i);

    Arrays.sort(nums);
    Arrays.sort(queryAndIndexes, (a, b) -> Integer.compare(a.getKey()[1], b.getKey()[1]));

    int i = 0; // nums' index
    for (Pair<int[], Integer> pair : queryAndIndexes) {
      final int[] query = pair.getKey();
      final int index = pair.getValue();
      final int x = query[0];
      final int m = query[1];
      while (i < nums.length && nums[i] <= m)
        bitTrie.insert(nums[i++]);
      if (i > 0 && nums[i - 1] <= m)
        ans[index] = bitTrie.getMaxXor(x);
    }

    return ans;
  }
}

class TrieNode:
  def __init__(self):
    self.children: List[Optional[TrieNode]] = [None] * 2


class BitTrie:
  def __init__(self, maxBit: int):
    self.maxBit = maxBit
    self.root = TrieNode()

  def insert(self, num: int) -> None:
    node = self.root
    for i in range(self.maxBit, -1, -1):
      bit = num >> i & 1
      if not node.children[bit]:
        node.children[bit] = TrieNode()
      node = node.children[bit]

  def getMaxXor(self, num: int) -> int:
    maxXor = 0
    node = self.root
    for i in range(self.maxBit, -1, -1):
      bit = num >> i & 1
      toggleBit = bit ^ 1
      if node.children[toggleBit]:
        maxXor = maxXor | 1 << i
        node = node.children[toggleBit]
      elif node.children[bit]:
        node = node.children[bit]
      else:  # Nothing in the Bit Trie.
        return 0
    return maxXor


class Solution:
  def maximizeXor(self, nums: List[int], queries: List[List[int]]) -> List[int]:
    ans = [-1] * len(queries)
    queryAndIndexes = sorted([(query, i) for i, query in enumerate(queries)],
                             key=lambda x: x[0][1])
    maxBit = int(math.log2(max(max(nums), max(x for x, _ in queries))))
    bitTrie = BitTrie(maxBit)

    nums.sort()

    i = 0  # nums' index
    for (x, m), index in queryAndIndexes:
      while i < len(nums) and nums[i] <= m:
        bitTrie.insert(nums[i])
        i += 1
      if i > 0 and nums[i - 1] <= m:
        ans[index] = bitTrie.getMaxXor(x)

    return ans