1707 Maximum XOR With an Element From Array

이번 문제는 Trie 구조를 이용해서 푸는 문제인데, Trie로 푸는 알고리즘 문제는 처음이었고,

그래서 풀 때, 시간이 좀 걸렸었다.

문제는 nums 배열에서 Queries 배열의 각 query의 두 번째 원소보다 같거나 작은 원소들과

첫 번째 원소와 XOR 연산을 했을 때의 최댓값을 찾는 문제이다.

 

어떠한 수 n과 임의의 수에 대해 XOR연산을 할 때, 만약 임의의 수가 5(101)와 4(100)라면, 공통되는 부분인 

10의 계산 결과를 저장할 수만 있다면, 2번 계산하지 않아도 될 것이다.

이 문제에서는 XOR연산의 이러한 점을 이용하여 비트 레벨에서 Trie구조를 만들어 문제를 풀 것이다.

 

일단 우리가 짤 Trie는 root가 상위 비트, child가 하위 비트인 Trie이다. 그래서 queries는 두 번째 원소가 낮은 원소인 query를 우선하여 정렬하고, nums 또한 정렬하여 추후 Trie를 구성할 때, Trie가 queries를 순회하면서 Trie에 저장된 원소의 크기가 점점 커질 수 있도록 했다. 그리고 추후 queries의 순서에 맞게 결과를 반환하기 위해서 각 query의 인덱스를 추가해주었다.

 

Trie를 구성할 때에는 int의 32비트에서 최우측 비트가 0번째라고 할 때, 부호 비트인 32번째 비트를 제외하고

최상위 비트인 31번째 비트부터 구성하였다. 그리고 가장 큰 XOR 연산의 결과를 찾을 때는 Trie에서 query의

첫 번째 원소의 최상위 비트부터 결과를 찾아나갔다.

 

C++

class Solution {
    class Trie{
        private:
            Trie* branch[2];
        public:
            Trie(){
                branch[0] = nullptr;
                branch[1] = nullptr;
            }
            static void build_trie(int num, Trie* root){
                Trie* curr = root;
                for(int i = 31; i >= 0; --i){
                    int value = (num >> i) & 1;
                    if (curr -> branch[value] == nullptr)
                        curr -> branch[value] = new Trie();
                    curr = curr -> branch[value];
                }
            }
            static int search(Trie* curr, int num){
                int res = 0;
                for(int i = 31; i >= 0  && curr != nullptr; --i){
                    int value = 1 - (num >> i) & 1;
                    if (curr -> branch[value] != nullptr){
                        res |= (1 << i);
                        curr = curr -> branch[value];
                    }
                    else {
                        curr = curr -> branch[1 - value];
                    }
                }
                return res;
            }
    };
public:
    vector<int> maximizeXor(vector<int>& nums, vector<vector<int>>& queries) {
        sort(nums.begin(), nums.end());
        for(size_t i = 0; i < queries.size(); ++i)
            queries[i].push_back(i);
        sort(queries.begin(), queries.end(), 
             [](vector<int>& a, vector<int>& b){return a[1] < b[1];});
        vector<int> res(queries.size(), -1);
        
        Trie* root = new Trie();
        int num_index = 0;
        for(size_t i = 0; i < queries.size(); ++i){
            while(num_index < nums.size() && nums[num_index] <= queries[i][1]){
                Trie::build_trie(nums[num_index], root);
                ++num_index;
            }
            if (num_index)
                res[queries[i][2]] = Trie::search(root, queries[i][0]);
        }
        
        return res;
    }
};

Python

class Solution:   
    def maximizeXor(self, nums: List[int], queries: List[List[int]]) -> List[int]:
        class Trie:
            def __init__(self):
                self.branch = {}
            def build_trie(self, num, root):
                curr = root.branch
                for i in range(31,-1,-1):
                    value = (num >> i) & 1
                    if value not in curr:
                        curr[value] = {}
                    curr = curr[value]
                
            def search(self, num, curr):
                res = 0
                if not self.branch:
                    return -1
                curr = self.branch
                res = 0
                for i in range(31,-1,-1):
                    value = 1 - (num >> i) & 1
                    if value in curr:
                        curr = curr[value]
                        res |= (1 << i)
                    else:
                        curr = curr[1 - value]
                return res
        
        res = [ -1 for i in range(len(queries))]
        nums.sort()
        for i in range(len(queries)):
            queries[i].append(i)
        queries.sort(key = lambda x : x[1])
        
        num_index = 0
        root = Trie()
        for i in range(len(queries)):
            while num_index < len(nums) and nums[num_index] <= queries[i][1]:
                root.build_trie(nums[num_index], root)
                num_index += 1
            if num_index:
                res[queries[i][2]] = root.search(queries[i][0], root)
                
        return res
        

처음에 C++코드를 파이썬으로 그대로 이식했는데 시간 초과가 떴다

그래서 Trie의 branch를 리스트에서 딕셔너리로 바꾸어 주었더니 또 잘 됐다.

흠... 리스트 오버헤드가 배열보다 앞도적으로 크거나 그러진 않을 텐데 뭐 때문이지..?