[2018 IUPC] 백준 15782 Calculate! 2

문제 링크

2018/07/26 추가내용.

참고) post-order대신 pre-order로 순회 해도 같은 효과를 볼 수 있으며, 코드를 조금 더 편하게, 더 직관적으로 짤 수 있다.

pre-order로 순회하는 코드는 여기를 참고하자.

1번 정점을 루트로 하는 트리가 주어지고, M개의 질의가 주어졌을 때 다음과 같은 명령을 실행한다.

• 1 x꼴로 주어지는 질의에는 정점 x를 포함한 x의 모든 자손들의 가중치를 전부 XOR한 값을 출력.
• 2 x y꼴로 주어지는 질의에는 정점 x를 포함한 x의 모든 자손들의 가중치에 각각 y를 XOR함 .

얼핏 봐선 그냥 하면 되는거 아닌가 싶겠지만 범위가 굉장히 크다.

정점의 수 N과 질의의 수 M에 대해

$(1\leq N \leq 100,000)$, $(1\leq M \leq 500,000)$ 

의 범위를 가진다.

질의가 주어질 때마다 정직하게 연산을 하게되면

이런 트리에서 쿼리 명령으로

"1 1"이나

"2 1 10"이 주어질 경우 한 번의 질의에 $O(N)$만큼 시간이 걸린다. M개의 쿼리가 모두 저런 식으로 주어졌을 때 $O(NM)$의 시간복잡도로 당연히 시간 초과를 받게 된다.

따라서 우리는 다른 방법을 찾아야 한다.

결론부터 말하면 세그먼트 트리를 이용하여 풀 수 있다. 세그먼트 트리를 모른다면 다음 링크에서 공부하고 오자.

백준 - 세그먼트 트리

백준 - 세그먼트 트리 lazy propagation

주어진 트리에서 어떤 부모와 그 모든 자손을 하나의 구간으로 볼 수 있기 때문에 세그먼트 트리로 해결할 수 있다. 하지만 이를 구현하기가 까다롭다.

트리의 모든 정점을 세그먼트 트리의 리프 노드로 만들되, 구간 업데이트를 적용하기 위해 자손과 부모를 연속하게 배치해야 한다. 또한 어떤 정점이 주어졌을 때 그 정점과 모든 자손이 세그먼트 트리의 어느 범위에 속하는지 알 수 있어야 한다.

문제의 예제 2번을 예로 들어 트리를 어떻게 세그먼트 트리로 만들 것인지 단계별로 살펴보자.

먼저 어떤 정점 $x$에 대해, $x$가 세그먼트 트리에서 몇 번째 리프노드인지 저장하는 $numbering$, $x$의 자식의 수가 몇 개인지 저장하는 $numChild$배열이 필요하다. 이 배열들의 쓰임새는 나중에 설명하겠다.

일단 정점을 재배치하기 위해 트리를 후위순회로 탐색한다. 이는 항상 자식이 부모보다 먼저 나열됨을 보장하기도 하면서, 자식의 수를 쉽게 구할 수 있는 방법이다.

$step\,0)$

초기상태이다. 

정점 번호	1	2	3	4	5	6	7
numbering
numChild

$step\,1)$

처음으로 방문하는 정점은 2번 정점이고 자식은 없다.

정점 번호	1	2	3	4	5	6	7
numbering		1
numChild		0

$step\,2)$

두 번째로 3번 정점을 방문한다. 자식은 없다.

정점 번호	1	2	3	4	5	6	7
numbering		1	2
numChild		0	0

$step\,3)$

세 번째로 5번 정점을 방문한다. 

정점 번호	1	2	3	4	5	6	7
numbering		1	2		3
numChild		0	0		0

$step\,4)$

네 번째로 7번 정점을 방문한다.

정점 번호	1	2	3	4	5	6	7
numbering		1	2		3		4
numChild		0	0		0		0

$step\,5)$

다섯 번째로 6번 정점을 방문하게 된다. 자식의 수는 1이다.

정점 번호	1	2	3	4	5	6	7
numbering		1	2		3	5	4
numChild		0	0		0	1	0

$step\,6)$

여섯 번째로 4번 정점을 방문하게 된다. 자식의 수는 3이다.

정점 번호	1	2	3	4	5	6	7
numbering		1	2	6	3	5	4
numChild		0	0	3	0	1	0

$step\,7)$

마지막으로 1번 정점을 방문하게 된다. 자식의 수는 6이다.

정점 번호	1	2	3	4	5	6	7
numbering	7	1	2	6	3	5	4
numChild	6	0	0	3	0	1	0

트리를 세그먼트 트리로 변환하기 위한 사전 준비가 끝났다. 위 과정은 코드상에서 init_seg_data(int)함수에 구현되어있다.

우선 $numbering$배열을 이용해 정점들의 순서를 재배열한 뒤(사실 코드상에선 그냥 위의 과정 중에 배열 p에다가 순서대로 저장한다) 세그먼트 트리를 만든다. 이 세그먼트 트리는 가중치를 저장하지만 아래 그림에선 나타내지 않도록 하겠다.

세그먼트 트리의 $i$번째 리프노드가 원래 트리의 $n$번 정점이라면, 이를 다음과 같이 나타내자.

세그먼트 트리를 만들어서 그려보면

이런 모습이다.

세그먼트 트리를 만들었으니 이제 쿼리를 해결하자. 어떤 정점 $x$가 주어지면 세그먼트 트리에서 그 정점의 위치 $i$는 $numbering[x]$임을 알고있다. 후위순회의 결과를 기반으로 정점들이 재배열되었기 때문에, 자신보다 자식이 항상 앞에 있으면서 연속적으로 배치되어 있다. 따라서 $x$의 모든 자손과 자기 자신을 포함하는 범위는 $[i - numChild[x], i]$이다. 이해가 안간다면 위 그림과 원래 트리의 그림을 비교해보자. 쉽게 알 수 있다.

쿼리문을 수행해야 하는 세그먼트 트리의 범위를 구했으므로 문제를 쉽게 풀 수 있게 된다.

주의할 점은 lazy update에서 변형이 조금 일어나는데, 어떤 수 $x$에 대해 $x\oplus x = 0$이므로 XOR연산을 $k$번 해야할 때 $k$가 짝수이면 연산하지 않아도 되고, $k$가 홀수이면 한 번만 XOR해주면 된다.

정답 코드

	#include <iostream>
	#include <vector>
	#include <algorithm>
	#include <cmath>
	using namespace std;
	vector<int> seg_tree, lazy, p(1), weight;
	vector<vector<int>> ori_tree;
	int n, m, numChild[100001], numbering[100001];
	bool check[100001];
	int init_seg_data(int now) {
	check[now] = true;
	int child = 0;
	for (int next : ori_tree[now]) {
	if(!check[next])
	child += init_seg_data(next);
	}
	numbering[now] = p.size();
	p.push_back(weight[now]);
	numChild[now] = child;
	return child + 1;
	}
	void init_seg_tree(int node, int start, int end) {
	if (start == end) {
	seg_tree[node] = p[start];
	}
	else {
	init_seg_tree(node * 2, start, (start + end) / 2);
	init_seg_tree(node * 2 + 1, (start + end) / 2 + 1, end);
	seg_tree[node] = seg_tree[node * 2] ^ seg_tree[node * 2 + 1];
	}
	}
	void update_lazy(int node, int start, int end) {
	if (lazy[node]) {
	if ((end - start + 1) & 1)
	seg_tree[node] ^= lazy[node];
	if (start != end) {
	lazy[node * 2] ^= lazy[node];
	lazy[node * 2 + 1] ^= lazy[node];
	}
	lazy[node] = 0;
	}
	}
	void update_range(int node, int start, int end, int i, int j, int diff) {
	update_lazy(node, start, end);
	if (start > j || end < i)
	return;
	else if (i <= start && end <= j) {
	if ((end - start + 1) & 1)
	seg_tree[node] ^= diff;
	if (start != end) {
	lazy[node * 2] ^= diff;
	lazy[node * 2 + 1] ^= diff;
	}
	return;
	}
	update_range(node * 2, start, (start + end) / 2, i, j, diff);
	update_range(node * 2 + 1, (start + end) / 2 + 1, end, i, j, diff);
	seg_tree[node] = seg_tree[node * 2] ^ seg_tree[node * 2 + 1];
	}
	int query(int node, int start, int end, int i, int j) {
	update_lazy(node, start, end);
	if (start > j || end < i)
	return 0;
	else if (i <= start && end <= j) {
	return seg_tree[node];
	}
	return query(node * 2, start, (start + end) / 2, i, j) ^ query(node * 2 + 1, (start + end) / 2 + 1, end, i, j);
	}
	int main() {
	ios_base::sync_with_stdio(false);
	cin.tie(0);
	cin >> n >> m;
	weight.resize(n + 1);
	ori_tree.resize(n + 1);
	int h = (int)ceil(log2(n));
	int tree_size = (1 << (h + 1));
	seg_tree.resize(tree_size);
	lazy.resize(tree_size);
	for (int i = 0; i < n - 1; i++) {
	int u, v;
	cin >> u >> v;
	ori_tree[u].push_back(v);
	ori_tree[v].push_back(u);
	}
	for (int i = 1; i < n + 1; i++) {
	cin >> weight[i];
	}
	init_seg_data(1);
	init_seg_tree(1, 1, n);
	while (m--) {
	int cmd, x, y;
	cin >> cmd;
	if (cmd == 1) { // x포함 자손들의 가중치의 XOR값 출력
	cin >> x;
	cout << query(1, 1, n, numbering[x] - numChild[x], numbering[x]) << '\n';
	}
	else if (cmd == 2) { // 구간업데이트, y를 XOR
	cin >> x >> y;
	update_range(1, 1, n, numbering[x] - numChild[x], numbering[x], y);
	}
	}
	}

view raw 15782.cpp hosted with ❤ by GitHub