BOJ 13445 - 부분 수열 XOR

BOJ 13445 - 부분 수열 XOR

Trie를 쓰는건 뻔하지만 풀이를 좀 생각해보자.

아이디어 - XOR Range Sum

$i=0 \to i=n-1$ 까지 순회해가며 각각 $i$가 부분 수열에서 마지막이 되는 수일 때 정답을 쿼리해와야한다.

$a_0$ 를 trie에 넣어두었다고 하자.

$a_1$을 볼 땐 $a_0$ 자체만으로 $a_0 \oplus a_1$ 이므로 신경써주지 않아도 된다.

그럼 이제 $a_1$ 을 Trie에 넣었다고 생각해보자.

$a_2$와 $\oplus$ 해서 $k$ 미만이 되게 하는 수를 그냥 세면 안된다.

왜냐면 trie에 들어있는 수들은 그냥 $a_0$ 과 $a_1$일 뿐이지 $a_0 \oplus a_1$ 과 $a_1$ 이 아니기 때문이다.

그러므로 $a_1$ 을 넣을 때 지금까지 있는 것들을 $a_1$ 과 $\oplus$ 한 것으로 바꿔주어야 한다는 것인데 이걸 어떻게 할 수 있을까?

우리가 지금까지 $i$ 를 거치며 트리에 삽입된 모든 수에 $\oplus$ 연산을 누적했다고 할 때, 그 누적된 수를 $X$ 라고 하자.

$\oplus$ 는 교환법칙과 결합법칙이 성립하기 때문에 가능한 일이다.

이진수 $X$의 $2^k$ 자리가 $1$이라는 것은 Trie에서 $2^k$ 위치를 가리키는 노드들은 두 자식을 swap해서 가지고 있어야 한다는 뜻이다.

실제로 매 단계마다 수를 집어넣어 Trie의 모든 정점에 대해서 Swap을 직접 해두게 할 수도 있지만 시간복잡도가 $O(N)$ 이라 시간안에 돌아가지 않는다.

$X$ 를 유지시키며 insert와 query 마다 $2^k$ 자리 일 때 $X$가 $2^k$자리가 1이라면, 두 자식 노드를 스위칭하게 만들면 문제가 해결된다.

아이디어 - Path less than K

현재 넣어야 할 수가 $x$ 이고 문제에서 제시되는 $K$ 가 있다고 하자.

Trie를 따라 내려가며 이미 이 path가 $K$ 보다 작아지는 경로인지를 계속 들고간다.

이미 작아졌다면 해당 Subtree의 모든 개수를 반환해준다.

그렇지 않다면 4가지 경우가 있다.

$x$와 $K$가 $2^k$ 자리가 $1$인지가 중요하다.

1. 모두 $1$일경우

이 땐, 현재 $2^k$ 가 $1$ 인 것들과 $\oplus$ 연산을 하면 이 자리가 $0$이 되므로 무조건 $K$ 이하가 되기 때문에 $1$ 을 의미하는 자식 노드의 모든 합을 더해준다.

$0$을 의미하는 자식 노드에서도 정답이 나올 수 있기 때문에 이쪽은 계속 재귀적으로 이동한다.

2. $x$ 는 $0$ $K$는 $1$일 경우

이 땐, 현재 $2^k$ 가 $0$ 인 것들과 $\oplus$ 연산을 하면 이 자리가 $0$이 되므로 무조건 $K$ 이하가 되기 때문에 $0$ 을 의미하는 자식 노드의 모든 합을 더해준다.

$1$을 의미하는 자식 노드에서도 정답이 나올 수 있기 때문에 이쪽은 계속 재귀적으로 이동한다.

3. $x$ 는 $0$ $K$는 $0$일 경우

$0$을 의미하는 자식 노드에서만 정답이 나올 수 있기 때문에 계속 재귀적으로 이동한다.

4. $x$ 는 $1$ $K$는 $0$일 경우

$1$을 의미하는 자식 노드에서만 정답이 나올 수 있기 때문에 계속 재귀적으로 이동한다.

이 모든건 아이디어 1에서 Swap과 병행되어야 한다는 점을 고려하자.

구현

아이디어는 어렵지 않게 떠올렸는데 오랜만에 XOR Trie를 짜니까 실수가 잦아 오래걸렸다.

int X = 0;  
  
struct trie {  
   int cnt_sum = 0, cnt = 0;  
   trie *zero = 0, *one = 0;  
   int insert(int x, int depth) {  
      if (depth == -1) {  
         cnt_sum++;  
         cnt++;  
         return cnt_sum;  
      }  
      if (X & (1 << depth))swap(one, zero);  
      if ((1 << depth) & x) {  
         if (!one) one = new trie();  
         one->insert(x, depth - 1);  
      } else {  
         if (!zero) zero = new trie();  
         zero->insert(x, depth - 1);  
      }  
      if (X & (1 << depth))swap(one, zero);  
      cnt_sum = (one ? one->cnt_sum : 0) + (zero ? zero->cnt_sum : 0);  
      return cnt_sum;  
   }  
  
   void print(int depth, string &s) {  
      if (cnt) {  
         cout << s << endl;  
      }  
      if (depth == -1) return;  
      if (X & (1 << depth))swap(one, zero);  
      s += '1';  
      if (one) one->print(depth - 1, s);  
      s.pop_back();  
      s += '0';  
      if (zero)zero->print(depth - 1, s);  
      s.pop_back();  
      if (X & (1 << depth))swap(one, zero);  
   }  
  
   int query(int x, int depth, int k, int less) {  
      if (depth == -1) return less ? cnt_sum : 0;  
  
      int x_bit = !!(x & (1 << depth));  
      int k_bit = !!(k & (1 << depth));  
  
      int ret = 0;  
      if (X & (1 << depth))swap(one, zero);  
      if (less) {  
         ret = cnt_sum;  
      } else if (k_bit == 1 && x_bit == 1) {  
         ret = (!zero ? 0 : zero->query(x, depth - 1, k, less)) + (!one ? 0 : one->query(x, depth - 1, k, 1));  
      } else if (k_bit == 1 && x_bit == 0) {  
         ret = (!zero ? 0 : zero->query(x, depth - 1, k, 1)) + (!one ? 0 : one->query(x, depth - 1, k, less));  
      } else if (k_bit == 0 && x_bit == 1) {  
         if (!one) ret = 0;  
         else ret = one->query(x, depth - 1, k, less);  
      } else if (k_bit == 0 && x_bit == 0) {  
         if (!zero) ret = 0;  
         else ret = zero->query(x, depth - 1, k, less);  
      }  
      if (X & (1 << depth))swap(one, zero);  
      return ret;  
   }  
};  
  
void solve() {  
   int n, k;  
   cin >> n >> k;  
   vi a(n);  
   fv(a);  
  
   trie t;  
   ll ans = 0;  
  
   for (int i: a) debug(bitset<5>(i));  
  
   int mx = 30;  
   for (int i = 0; i < n; i++) {  
      int cur = a[i];  
      if (cur < k) ans++;  
      ans += t.query(cur, mx, k, 0);  
      X ^= cur;  
      t.insert(cur, mx);  
   }  
   cout << ans;  
}