[C++] 백준 16903번: 수열과 쿼리 20

문제 링크

16903번: 수열과 쿼리 20

문제 요약

0이 하나 포함된 배열 A에 대해서 다음의 쿼리들을 구현해야 한다.
1 x: A에 x를 추가한다.
2 x: A에서 x를 제거한다. A에 x가 두 개 이상 있는 경우에는 하나만 삭제한다. 항상 A에 x가 있는 쿼리만 주어진다.
3 x: A에 포함된 각각의 원소와 x를 XOR 연산을 한 다음, 가장 큰 값을 출력한다.

접근 방법

1번 쿼리에 대해서는 x를 비트셋으로 바꿔서 트라이에 저장해줍니다. 트라이의 각 노드는 0과 1 두 비트를 가리키는 자식 노드에 대한 포인터 변수를 가지고 있습니다. 또한, 현재 노드에 해당하는 수를 저장할 cnt 변수를 하나 가지고 있습니다. 트라이에서 비트셋을 따라가면서 필요에 따라 노드들을 생성해주고, 방문 할때마다 cnt 변수를 증가시켜 줍니다.

제거시에는 반대로 하면 됩니다. 일단 cnt를 감소시켜주고, 만약 0이 된 노드가 있다면 delete 연산자를 통해서 제거해줍니다. 그리고 제거된 노드의 부모 노드에서는 포인터 변수에 NULL을 대입해줍니다.

이제 3번 쿼리입니다. XOR은 두 비트가 서로 달라야 1이 됩니다. XOR한 값이 최대한 커지기 위해서는 가능한 한 대응되는 두 비트가 서로 달라야 합니다. x를 비트셋으로 바꿔주고, 트라이의 루트부터 x의 현재 비트와 다른 비트에 해당하는 자식 노드가 존재한다면 그쪽으로 이동합니다. 만약 그렇지 않다면 동일한 비트의 노드로 이동합니다. 그렇게해서 리프에 도달하게 되면 두 수의 XOR한 값을 리턴해주고 출력하면 됩니다.

0이 이미 배열에 들어있다는 점에 주의를 해야 합니다.

풀고 나서 난이도 기여를 보니 세그먼트 트리로도 풀린다고 하는데, 세그먼트 트리로 푸는 방법은 아직 생각이 나진 않습니다. 생각해낸다면 여기에 추가해볼 생각입니다.

코드

#include <bits/stdc++.h>

using namespace std;

struct Trie {
	int cnt;
	Trie* node[2];

	Trie() {
		cnt = 0;
		node[0] = node[1] = NULL;
	}

	void insert(bitset<32>& bit, int idx) {
		cnt++;

		if (idx == -1)
			return;

		if (node[bit[idx]] == NULL) {
			Trie* trie = new Trie();
			node[bit[idx]] = trie;
		}

		node[bit[idx]]->insert(bit, idx - 1);
	}

	void erase(bitset<32>& bit, int idx) {
		cnt--;

		if (idx == -1)
			return;
		
		node[bit[idx]]->erase(bit, idx - 1);

		if (!node[bit[idx]]->cnt) {
			delete node[bit[idx]];
			node[bit[idx]] = NULL;
		}
	}

	int find(bitset<32>& bit, bitset<32>& res, int idx) {
		if (idx == -1)
			return (int)res.to_ulong();

		if (node[!bit[idx]] != NULL) {
			res[idx] = 1;
			return node[!bit[idx]]->find(bit, res, idx - 1);
		}
		else {
			res[idx] = 0;
			return node[bit[idx]]->find(bit, res, idx - 1);
		}
	}
};

int main(void) {
	ios::sync_with_stdio(false);
	cin.tie(0);
	cout.tie(0);

	int m;
	cin >> m;

	Trie* root = new Trie();
	bitset<32> zero(0);
	root->insert(zero, 31);

	while (m--) {
		int q, x;
		cin >> q >> x;

		bitset<32> tmp(x);
		bitset<32> res(x);

		switch (q) {
		case 1:
			root->insert(tmp, 31);
			break;
		case 2:
			root->erase(tmp, 31);
			break;
		case 3:
			cout << root->find(tmp, res, 31) << '\n';
		}
	}

	delete root;
	return 0;
}