Maximum XOR With an Element From Array LeetCode solution in Python

1
class Trie:
2
    def __init__(self):
3
        self.root = {}
4

5
    def insert(self, num):
6
        p = self.root
7
        for i in range(31, -1, -1):
8
            cur = (num >> i) & 1
9
            if cur not in p:
10
                p[cur] = {}
11
            p = p[cur]
12

13
    def query(self, num):
14
        if not self.root:
15
            return -1
16
        p, ans = self.root, 0
17
        for i in range(31, -1, -1):
18
            cur = (num >> i) & 1
19
            if 1 - cur in p:
20
                p = p[1 - cur]
21
                ans |= 1 << i
22
            else:
23
                p = p[cur]
24
        return ans
25

26

27
class Solution:
28
    def maximizeXor(self, nums: List[int], queries: List[List[int]]) -> List[int]:
29
        nums.sort()
30
        queries = sorted(enumerate(queries), key=lambda x: x[1][1])
31
        trie = Trie()
32
        ans = [-1] * len(queries)
33
        j = 0
34
        for i, (x, m) in queries:
35
            while j < len(nums) and nums[j] <= m:
36
                trie.insert(nums[j])
37
                j += 1
38
            ans[i] = trie.query(x)
39
        return ans