421. Maximum XOR of Two Numbers in an Array

Approach 1: Hash Set

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

C++JavaPython

class Solution {
 public:
  int findMaximumXOR(vector<int>& nums) {
    const int maxNum = ranges::max(nums);
    if (maxNum == 0)
      return 0;
    const int maxBit = static_cast<int>(log2(maxNum));
    int ans = 0;
    int mask = 0;

    // If ans is 11100 when i = 2, it means that before we reach the last two
    // bits, 11100 is the maximum XOR we have, and we're going to explore if we
    // can get another two '1's and put them into ans.
    for (int i = maxBit; i >= 0; --i) {
      // Mask grows like: 100...000, 110...000, 111...000, ..., 111...111.
      mask |= 1 << i;
      unordered_set<int> prefixes;
      // We only care about the left parts,
      // If i = 2, nums = {1110, 1011, 0111}
      //    -> prefixes = {1100, 1000, 0100}
      for (const int num : nums)
        prefixes.insert(num & mask);
      // If i = 1 and before this iteration, the ans is 10100, it means that we
      // want to grow the ans to 10100 | 1 << 1 = 10110 and we're looking for
      // XOR of two prefixes = candidate.
      const int candidate = ans | 1 << i;
      for (const int prefix : prefixes)
        if (prefixes.count(prefix ^ candidate)) {
          ans = candidate;
          break;
        }
    }

    return ans;
  }
};

class Solution {
  public int findMaximumXOR(int[] nums) {
    final int maxNum = Arrays.stream(nums).max().getAsInt();
    if (maxNum == 0)
      return 0;
    final int maxBit = (int) (Math.log(maxNum) / Math.log(2));
    int ans = 0;
    int mask = 0;

    // If ans is 11100 when i = 2, it means that before we reach the last two
    // bits, 11100 is the maximum XOR we have, and we're going to explore if we
    // can get another two '1's and put them into ans.
    for (int i = maxBit; i >= 0; --i) {
      // Mask grows like: 100...000, 110...000, 111...000, ..., 111...111.
      mask |= 1 << i;
      Set<Integer> prefixes = new HashSet<>();
      // We only care about the left parts,
      // If i = 2, nums = {1110, 1011, 0111}
      //    . prefixes = {1100, 1000, 0100}
      for (final int num : nums)
        prefixes.add(num & mask);
      // If i = 1 and before this iteration, the ans is 10100, it means that we
      // want to grow the ans to 10100 | 1 << 1 = 10110 and we're looking for
      // XOR of two prefixes = candidate.
      final int candidate = ans | 1 << i;
      for (final int prefix : prefixes)
        if (prefixes.contains(prefix ^ candidate)) {
          ans = candidate;
          break;
        }
    }

    return ans;
  }
}

class Solution:
  def findMaximumXOR(self, nums: List[int]) -> int:
    maxNum = max(nums)
    if maxNum == 0:
      return 0
    maxBit = int(math.log2(maxNum))
    ans = 0
    mask = 0

    # If ans is 11100 when i = 2, it means that before we reach the last two
    # bits, 11100 is the maximum XOR we have, and we're going to explore if we
    # can get another two '1's and put them into ans.
    for i in range(maxBit, -1, -1):
      # Mask grows like: 100...000, 110...000, 111...000, ..., 111...111.
      mask |= 1 << i
      # We only care about the left parts,
      # If i = 2, nums = [1110, 1011, 0111]
      #    -> prefixes = [1100, 1000, 0100]
      prefixes = set([num & mask for num in nums])
      # If i = 1 and before this iteration, the ans is 10100, it means that we
      # want to grow the ans to 10100 | 1 << 1 = 10110 and we're looking for
      # XOR of two prefixes = candidate.
      candidate = ans | 1 << i
      for prefix in prefixes:
        if prefix ^ candidate in prefixes:
          ans = candidate
          break

    return ans

Approach 2: Bit Trie

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

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

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

class Solution {
 public:
  int findMaximumXOR(vector<int>& nums) {
    const int maxNum = ranges::max(nums);
    if (maxNum == 0)
      return 0;
    const int maxBit = static_cast<int>(log2(maxNum));
    int ans = 0;
    BitTrie bitTrie(maxBit);

    for (const int num : nums) {
      ans = max(ans, bitTrie.getMaxXor(num));
      bitTrie.insert(num);
    }

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

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

class Solution {
  public int findMaximumXOR(int[] nums) {
    final int maxNum = Arrays.stream(nums).max().getAsInt();
    if (maxNum == 0)
      return 0;
    final int maxBit = (int) (Math.log(maxNum) / Math.log(2));
    int ans = 0;
    BitTrie bitTrie = new BitTrie(maxBit);

    for (final int num : nums) {
      ans = Math.max(ans, bitTrie.getMaxXor(num));
      bitTrie.insert(num);
    }

    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:  # There's nothing the Bit Trie.
        return 0
    return maxXor


class Solution:
  def findMaximumXOR(self, nums: List[int]) -> int:
    maxNum = max(nums)
    if maxNum == 0:
      return 0
    maxBit = int(math.log2(maxNum))
    ans = 0
    bitTrie = BitTrie(maxBit)

    for num in nums:
      ans = max(ans, bitTrie.getMaxXor(num))
      bitTrie.insert(num)

    return ans