Search code examples
algorithmdata-structuresmaxxor

LeetCode 1707. 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. This python solution uses Trie, but still LeetCode shows TLE?

 import operator

 class TrieNode:
      def __init__(self):
          self.left=None
          self.right=None
   
 class Solution:
     def insert(self,head,x):
         curr=head
         for i in range(31,-1,-1):
             val = (x>>i) & 1
             if val==0:
                if not curr.left:
                   curr.left=TrieNode()
                       curr=curr.left
                   else:
                       curr=curr.left
               else:
                   if not curr.right:
                       curr.right=TrieNode()
                       curr=curr.right
                   else:
                       curr=curr.right
       
           
       def maximizeXor(self, nums: List[int], queries: List[List[int]]) -> List[int]:
           res=[-10]*len(queries)
           nums.sort()
           for i in range(len(queries)):
               queries[i].append(i)
           queries.sort(key=operator.itemgetter(1))
           head=TrieNode()
          
           for li in queries:
               max=0
               xi,mi,index=li[0],li[1],li[2]
               m=2**31
               node = head
               pos=0
               if mi<nums[0]:
                   res[index]=-1
                   continue
               for i in range(pos,len(nums)):
                   if mi<nums[i]:
                       pos=i
                       break
                   self.insert(node,nums[i])
               node=head
               for i in range(31,-1,-1):
                   val=(xi>>i)&1
                   if val==0:
                       if node.right:
                           max+=m
                           node=node.right
                       else:
                           node=node.left
                   else:
                       if node.left:
                           max+=m
                           node=node.left
                       else:
                           node=node.right
                   m>>=1
               res[index]=max
           return -1

Solution

  • here is alternative Trie implement to solve this problem:

    [Notes: 1) max(x XOR y for y in A); 2) do the greedy on MSB bit; 3) sort the queries]

    class Trie:
        def __init__(self):
            self.root = {}
        
        def add(self, n):
            p = self.root
            for bitpos in range(31, -1, -1):
                bit = (n >> bitpos) & 1
                if bit not in p:
                    p[bit] = {}
                p = p[bit]
        
        def query(self, n):
            p = self.root
            ret = 0
            if not p:
                return -1
            for bitpos in range(31, -1, -1):
                bit = (n >> bitpos) & 1
                inverse = 1 - bit
                if inverse in p:
                    p = p[inverse]
                    ret |= (1 << bitpos)
                else:
                    p = p[bit]
                    
            return ret
    
    class Solution:
        def maximizeXor(self, nums: List[int], queries: List[List[int]]) -> List[int]:
            
            n = len(nums)
            trie = Trie()
            q = sorted(enumerate(queries), key = lambda x: x[1][1])
            nums.sort()
            res = [-1] * len(queries)
            i = 0
            for index, (x, m) in q:
                while i < n and nums[i] <= m:
                    trie.add(nums[i])
                    i += 1
                res[index] = trie.query(x)
            return res