LeetCode 1707. Maximum XOR With an Element From Array - algorithm

Question

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

Answer 1

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

Answer 2

The problem is that you're building a fresh Trie for each query. And to make matters worse, use linear search to find the maximum value <= mi in nums. You'd be better off by simply using
max((n for n in nums if n <= mi), key=lambda n: n^xi, default=-1)
The solution here would be to build the trie right at the start and simply filter for values smaller than mi using that trie:
import math
import bisect
def dump(t, indent=''):
if t is not None:
print(indent, "bit=", t.bit, "val=", t.val, "lower=", t.lower)
dump(t.left, indent + '\tl')
dump(t.right, indent + '\tr')
class Trie:
def __init__(self, bit, val, lower):
self.bit = bit
self.val = val
self.lower = lower
self.left = None
self.right = None
def solve(self, mi, xi):
print('-------------------------------------------')
print(self.bit, "mi(b)=", (mi >> self.bit) & 1, "xi(b)=", (xi >> self.bit) & 1, "mi=", mi, "xi=", xi)
dump(self)
if self.val is not None:
# reached a leave of the trie => found matching value
print("Leaf")
return self.val
if mi & (1 << self.bit) == 0:
# the maximum has a zero-bit at this position => all values in the right subtree are > mi
print("Left forced by max")
return -1 if self.left is None else self.left.solve(mi, xi)
# pick based on xor-value if possible
if (xi >> self.bit) & 1 == 0 and self.right is not None and (mi > self.right.lower or mi == ~0):
print("Right preferred by xi")
return self.right.solve(mi, xi)
elif (xi >> self.bit) & 1 == 1 and self.left is not None:
print("Left preferred by xi")
return self.left.solve(~0, xi)
# pick whichever is available
if self.right is not None and (mi > self.right.lower or mi == ~0):
print("Only right available")
return self.right.solve(mi, xi)
elif self.left is not None:
print("Only left available")
return self.left.solve(~0, xi)
else:
print("None available")
return -1
def build_trie(nums):
nums.sort()
# msb of max(nums)
max_bit = int(math.log(nums[-1], 2)) # I'll just assume that nums is never empty
print(max_bit)
def node(start, end, bit, template):
print(start, end, bit, template, nums[start:end])
if end - start == 1:
# reached a leaf
return Trie(0, nums[start], nums[start])
elif start == end:
# a partition without values => no Trie-node
return None
# find pivot for partitioning based on bit-value of specified position (bit)
part = bisect.bisect_left(nums, template | (1 << bit), start, end)
print(part)
# build nodes for paritioning
res = Trie(bit, None, nums[start])
res.left = node(start, part, bit - 1, template)
res.right = node(part, end, bit - 1, template | (1 << bit))
return res
return node(0, len(nums), max_bit, 0)
class Solution:
def maximizeXor(self, nums: List[int], queries: List[List[int]]) -> List[int]:
trie = build_trie(nums)
return [trie.solve(mi if mi <= nums[-1] else ~0, xi) for xi, mi in queries]
I've been a bit lazy and simply used ~0 to signify that the maximum can be ignored since all values in the subtree are smaller than mi. The basic idea is that ~0 & x == x is true for any integer x. Not quite as simple as #DanielHao's answer, but capable of handling streams of queries.