[C++] Trie로 최대/최소 XOR 문제 풀기(XOR Trie)

 제가 코드포스 블루를 달았을 때 사용했던 알고리즘이기도 하고, 이걸 변형해서 제가 문제도 냈었기 때문에 한 번 설명해보고자 합니다. (정확한 설명은 아닐 수 있지만, 최대한 직관적이게 적어볼려고 하니 컨셉만 잘 전달되었으면 합니다.)

 

Trie(트라이)는 문자열을 트리 형태로 저장하는 자료구조입니다.  아래 사진을 보면 이해하기 쉬우실 겁니다.

A trie for keys "A", "to", "tea", "ted", "ten", "i", "in", and "inn".  (출처 : https://en.wikipedia.org/wiki/Trie)

 

보면은 해당 트라이에 삽입된 단어들은 "A", "to", "tea", "ted", "ten", "i", "in", and "inn" 인데, prefix가 겹친다면 같은 노드를 타고 내려가도록 트리가 생성되어 있습니다. 트라이의 특징으로는 Insertion / Deletion / Search 하는데 전부 O(|S|) 만에 가능합니다. 

 

그런데 이러한 자료구조의 특징을 통해서 우리는최대/최소 XOR 문제들도 풀 수 있습니다. 이런 XOR 문제를 풀 때 사용하는 Trie를 편의상 XOR Trie라고 부르겠습니다. 그러면 이제 문제를 예시로 하나 들어 이 문제를 해결해가면서 XOR Trie를 설명하도록 하겠습니다.

BOJ 13505 : 두 수 XOR - https://www.acmicpc.net/problem/13505

13505번: 두 수 XOR

문제에서 요구하는 바는 수열이 주어져 있을 때 두 수의 XOR이 제일 큰 결과값이 무엇인지를 물어보고 있습니다. N의 범위가 작다면 쉽게 O(N^2)으로 브루트포싱해서 풀 수는 있겠지만 N의 범위가 너무커서 시간내로 돌아가지 않습니다. 이 문제를 어떻게 Trie로 풀 수 있을까요?

 

먼저, XOR 연산의 특징을 생각해봅시다. XOR는 같은 자리의 비트끼리만 연산을 합니다. 자리가 다르다면 연산의 결과가 서로 독립적으로 계산이 되어 영향을 서로주지 않죠. 이러한 특징은 이 문제를 푸는데 키가 됩니다. 왜냐하면, 특정 수와의 XOR한 결과가 제일 커지는 값은 최상위 비트부터 1이 되는지를 판단해서 찾아볼 수 있기 때문이죠.

 

어느 수 X가 있다고 가정합시다. 그리고 두 수 Y, Z가 있다고 합시다. 그리고 XOR한 값의 결과가 아래와 같다고 합시다. (1번째 비트를 제일 큰 비트라고 표현합시다. 즉, k번째 비트가 k+1번째 비트보다 표현하는 값이 큰 걸로 해석해주세요.)

• X ^ Y = k 번째 비트만 1로 설정됨.
• X ^ Z = k+1번째 비트부터 전부 1로 설정됨.

 

그렇다면 어느 수가 X와 XOR 했을 때 더 값이 클까요? 당연히 Y입니다. 왜냐면 k번째 비트가 1로 켜진 결과가 k번째 비트 아래의 비트들이 전부 1로 켜진 결과보다 더 크기 때문입니다. 상위 비트가 1로 켜진다면 그 하위의 비트들이 아무리 1로 많이 켜진다 해도 해당 값을 넘을 수가 없습니다. 그래서 이 문제는 좀 더 그리디한 접근으로 찾을 수 있습니다. 특정 값과 XOR를 한 결과가 상위 비트에서 1이 되는 수들과 0으로 되는 수들이 있으면 1로 켜지는 값을 우선적으로 찾아보는 겁니다. 그리고 그 하위 비트에서도 1이 되는 수들을 최대한 찾아보고요.

 

위 과정을 Trie를 활용해서 한다면 검색이 O(logA_i)만에 가능해집니다. 왜냐면 우리는 이진법으로 각 숫자를 표현해서 상위 비트부터 순서대로 Trie에 저장할 것이기 때문입니다. 근데 주의해야할 점은 비트를 몇 개 볼지는 항상 통일해야합니다. 그래야 제대로 비교가 가능하니깐요. 나올 수 있는 수들중에 제일 큰 수가 10^9이므로 대강 깊이를 40이라고 생각해서 구해봅시다. (즉 40개의 비트들을 전부 저장할 겁니다.) 

 

일단 Trie를 위한 구조체부터 먼저 선언하겠습니다. 그리고 각 수를 처리할 수 있는 전역변수들도 같이 선언해보죠.

deque<ll> bi;
ll val = 0;
// 전역 변수
// bi는 각 수를 이진법으로 표현한 값들을 저장하기 위한 배열
// val은 특정 값과의 XOR한 값 중 최댓값을 저장할 용도의 변수

struct node{
    node *childs[2];
    node(){
        childs[0] = nullptr;
        childs[1] = nullptr;
    }
}; // Trie를 위한 node 구조체

 

그리고 Insertion도 같이 구현해봅시다. 그 전에 우리는 각 숫자들을 이진수로 만들어서 저장하는 함수도 만들겁니다.

void add_n(ll n){
    bi.clear();
    while(n > 0){
        bi.push_front(n%2);
        n/=2;
    }
    while(bi.size() < 40)
        bi.push_front(0);
} // 숫자들을 이진법으로 표현해서 bi 배열에 넣어주는 함수, 2^39에 해당되는 비트까지 본다.

void add(node * tree, int idx){
    if(idx == 40) return; // 최대 깊이에 들어가면 더 이상 X
    ll x = bi[idx]; // 해당 비트의 값
    if(!tree -> childs[x]) // 한 번도 해당 비트에서 x값이 나온 적이 없었으면 생성 
        tree->childs[x] = new node();
    add(tree->childs[x], idx+1); // 해당 비트에 대한 노드에서 다음 비트값 탐색
} // Trie Insertion

 

그리고 이제 최댓값을 구하는 함수도 만들어봅시다. 만약에 특정 비트 위치에서 XOR 결과가 1이 될려면 두 비트의 값은 서로 달라야 합니다. 그래서 우리가 기준으로 잡은 값의 특정 비트에서 서로 다른 값인 수들이 존재한다면 해당 수들을 찾으면 되는 거고, 기준으로 잡은 값과 XOR한 결과들 중 최댓값은 해당 비트에 1이 들어오게 됩니다. 이 과정을 최상위 비트부터 찾아보면 최상위 비트들이 최대한 1이 되는 방향으로 찾게됩니다. (아마 코드로 이해하시는게 더 편할 것 같긴 합니다.)

void get(node* tree, int idx){
    if(idx == 40) return; // 최대 깊이까지 왔으면 그만.
    int x = bi[idx]; // 기준이 되는 값의 idx번째 비트의 값 (1 아니면 0)
    if(!tree->childs[x]){ // 만약에 반대인 값이 존재한다면
        val+=(1LL<<(long long)(39-idx)); // 최댓값에 해당 비트는 1이 들어오는 것을 명시
        get(tree->childs[1-x], idx+1); // 그리고 비트가 반대인 값들 중에서 찾음.
    }
    else{
        get(tree->childs[x], idx+1); // 아니면 최댓값에는 해당 비트가 0일꺼고, 그냥 같은 값들 중에서 찾음.
    }
}

 

위 함수들을 전부 합쳐서 예제 문제를 풀어보면 아래와 같은 코드로 풀 수 있습니다.

#include<bits/stdc++.h>
#define ll long long
using namespace std;
// Default

//////////////////// XOR TRIE ////////////////////

deque<ll> bi;
ll val = 0;

struct node{
    node *childs[2];
    node(){
        childs[0] = nullptr;
        childs[1] = nullptr;
    }
};

void add(node * tree, int idx){
    if(idx == 40) return;
    ll x = bi[idx];
    if(!tree -> childs[x])
        tree->childs[x] = new node();
    add(tree->childs[x], idx+1);
}

void add_n(ll n){
    bi.clear();
    while(n > 0){
        bi.push_front(n%2);
        n/=2;
    }
    while(bi.size() < 40)
        bi.push_front(0);
}

void get(node* tree, int idx){
    if(idx == 40) return;
    int x = bi[idx];
    if(!tree->childs[1-x]){
        get(tree->childs[x], idx+1);
    }
    else{
        val+=(1LL<<(long long)(39-idx));
        get(tree->childs[1-x], idx+1);
    }
}

//////////////////// XOR TRIE ////////////////////

int main(){
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    
    ll N;
    cin >> N;
    vector<ll> v(N); 
    node * root = new node();
    for(int i =0; i<N; i++){
        cin >> v[i];
        add_n(v[i]);
        add(root, 0);
    } // 각 값을 입력 받고 이를 이진수로 표현해서 Trie에 저장.
    
    ll maxx = 0; 
    for(int i =0; i<N; i++){
        val = 0; // 값 초기화
        add_n(v[i]); // 각 값을 기준으로
        get(root , 0); // XOR 최댓값을 찾음 (val이라는 변수에 저장됨.)
        maxx = max(maxx , val); // 해당 결과를 다른 값들을 기준으로 삼을 때의 XOR 최댓값과 비교
    }

    cout << maxx;

    return 0;
} // O(40N)

 

이렇게 문제를 풀게되면 O(40*N)만에 풀 수 있습니다. (별개로 이 문제의 다른 해로는 이분 탐색이 있긴 합니다. 그런데 사용되는 아이디어는 동일해요.)

 

이런식으로 Trie를 사용해서 XOR 문제도 풀 수 있습니다. 제 설명이 뭔가 정확한 증명을 토대로 하는게 아니라 직관적으로 이해가 되도록 최대한 설명을 했는데, 이게 제대로 전달이 되었을 지 모르겠습니다. (정확한 증명은 이 글에서 잘 설명이 되있는 것 같습니다.) 그래도 문제를 푸는 아이디어라도 어느정도 전달되었으면 좋겠습니다. 이제 아래에 추천 문제들을 남기면서 글을 끝내도록 하겠습니다. 다음에 다뤄볼 알고리즘은 뭐가 될지는 저도 궁금하네요.

 

추천문제

• BOJ 13505 : 두 수 XOR - 위에서 언급한 연습 문제입니다.
• BOJ 13504 : XOR 합 - BOJ 13505과 대표적인 XOR Trie 문제입니다. 구간 XOR를 활용해서 쉽게 풀 수가 있어요.
• BOJ 20919 : XOR 자료구조 - XOR Trie인데 삭제 등등도 구현해야합니다. 또한 XOR 결과가 최소가 되는 값도 찾아야 합니다.
• BOJ 30210 : BX 내기 - 제가 만들었던 문제인데, 위 아이디어를 십진수에 적용해서도 풀 수 있습니다. 한 번 풀어보세요 :)
• CP Round 882(Div.2) C번 : Vampiric Powers, anyone? - 제가 처음으로 XOR Trie로 풀었던 문제인데, 정해는 아닙니다. (대회 끝나고보니 좀 더 쉬운 정해가...)
• CP Round 888(Div.3) F번 : Lisa and the Martians - 좀 관찰이 필요한 XOR Trie 문제입니다.