Problem: Maximum XOR With an Element From Array

You are given an array `nums` consisting of non-negative integers. You are
also given a `queries` array, where `queries[i] = [xi, mi]`.

The answer to the `ith` query is the maximum bitwise `XOR` value of `xi` and
any element of `nums` that does not exceed `mi`. In other words, the answer is
`max(nums[j] XOR xi)` for all `j` such that `nums[j] <= mi`. If all elements
in `nums` are larger than `mi`, then the answer is `-1`.

Return an integer array `answer` where `answer.length == queries.length` and
`answer[i]` is the answer to the `ith` query.


Python solution:
import math
import itertools
import bisect
import collections
import string
import heapq
import functools
import sortedcontainers
from typing import List, Dict, Tuple, Iterator, Optional

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:
        return 0
    return maxXor


class IndexedQuery:
  def __init__(self, queryIndex: int, x: int, m: int):
    self.queryIndex = queryIndex
    self.x = x
    self.m = m

  def __iter__(self):
    yield self.queryIndex
    yield self.x
    yield self.m


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

    nums.sort()

    i = 0
    for queryIndex, x, m in sorted([IndexedQuery(i, x, m) for i, (x, m) in enumerate(queries)], key=lambda iq: iq.m):
      while i < len(nums) and nums[i] <= m:
        bitTrie.insert(nums[i])
        i += 1
      if i > 0 and nums[i - 1] <= m:
        ans[queryIndex] = bitTrie.getMaxXor(x)

    return ans


Java solution:
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 { // There's 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);
    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);

    Arrays.sort(nums);

    int i = 0; // nums' index
    for (IndexedQuery indexedQuery : getIndexedQueries(queries)) {
      final int queryIndex = indexedQuery.queryIndex;
      final int x = indexedQuery.x;
      final int m = indexedQuery.m;
      while (i < nums.length && nums[i] <= m)
        bitTrie.insert(nums[i++]);
      if (i > 0 && nums[i - 1] <= m)
        ans[queryIndex] = bitTrie.getMaxXor(x);
    }

    return ans;
  }

  private record IndexedQuery(int queryIndex, int x, int m){};

  private IndexedQuery[] getIndexedQueries(int[][] queries) {
    IndexedQuery[] indexedQueries = new IndexedQuery[queries.length];
    for (int i = 0; i < queries.length; ++i)
      indexedQueries[i] = new IndexedQuery(i, queries[i][0], queries[i][1]);
    Arrays.sort(indexedQueries, Comparator.comparingInt(IndexedQuery::m));
    return indexedQueries;
  }
}


C++ solution:
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 {  // There's nothing in the Bit Trie.
        return 0;
      }
    }
    return maxXor;
  }

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

struct IndexedQuery {
  int queryIndex;
  int x;
  int m;
};

class Solution {
 public:
  vector<int> maximizeXor(vector<int>& nums, vector<vector<int>>& queries) {
    vector<int> ans(queries.size(), -1);
    const int maxNumInNums = ranges::max(nums);
    const int maxNumInQuery = ranges::max_element(queries, ranges::less{},
                                                  [](const vector<int>& query) {
      return query[0];
    })->at(0);
    const int maxBit = static_cast<int>(log2(max(maxNumInNums, maxNumInQuery)));
    BitTrie bitTrie(maxBit);

    ranges::sort(nums);

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

    return ans;
  }

 private:
  vector<IndexedQuery> getIndexedQueries(const vector<vector<int>>& queries) {
    vector<IndexedQuery> indexedQueries;
    for (int i = 0; i < queries.size(); ++i)
      indexedQueries.emplace_back(i, queries[i][0], queries[i][1]);
    ranges::sort(
        indexedQueries, ranges::less{},
        [](const IndexedQuery& indexedQuery) { return indexedQuery.m; });
    return indexedQueries;
  }
};