취미로 코딩하기

※ 글쓴이는 취미로 코딩을 익혀보는 사람이라 정확하지 않은 내용을 담고 있을 수 있다 ※

이번에 볼 문제는 백준 18260번 문제인 Bessie's Snow Cow이다.
문제는 아래 링크를 확인하자.

https://www.acmicpc.net/problem/18260 

이 문제에서는 1번 노드를 루트로 갖는 트리에서 (1) x번 노드를 루트로 하는 서브트리에 색 c를 입히거나(단, 기존에 색을 가진 노드에 색을 칠해도 기존 색이 지워지지 않음) (2) x번 노드를 루트로 하는 서브트리에 속한 각 노드가 가지는 색의 개수의 합을 구하는 쿼리를 처리하는 문제이다.

이 문제에서는 서브트리에 대한 쿼리만을 처리할 것이므로 먼저 Euler Tour Technique을 이용해 주어진 트리를 배열로 펴 l구간합 lazy 세그먼트트리를 이용할 수 있게끔 인덱스를 새로 부여하는 전처리를 하자. 이 세그먼트트리는 각 노드 리프가 갖는 값이 그 노드가 갖는 색의 가짓수가 되게끔 관리할 것이고, 이렇게 관리된 세그먼트트리를 이용하면 (2)번 쿼리를 구간합 쿼리를 이용해 해결할 수 있음을 관찰할 수 있다. 이제 (1)번 쿼리를 처리하는 방법을 생각해보자.

만약 x번 노드를 루트로 하는 서브트리에 색이 c인 노드를 가진 노드가 전혀 없었다면 x번 노드를 루트로 갖는 서브트리의 노드들에 색의 종류를 전부 1씩 더하는 것으로 위의 목표를 달성할 수 있게 될 것이다. 그렇지 않은 경우는 다음과 같은 두 경우로 나누어 생각할 수 있다.

i) x번 노드에 c 색이 칠해져있지 않지만, 서브트리를 구성하는 노드 중 c로 이미 칠해진 노드가 있는 경우

이 경우 x번 노드의 서브트리에서 진행된 "c 칠하기" 시행을 "취소"한다면 위의 간단한 경우와 같이 쿼리를 처리할 수 있게 될 것이다. 이를 위해 기존에 입력으로 들어온 "노드 k에 c 칠하기" 시행들을 set 등을 이용해 관리하고 x번 노드의 서브트리에 속하는 모든 k에 대해 위에서 진행한 시행을 역으로 실행(k번 노드를 루트로 갖는 서브트리의 노드들에 색의 종류를 전부 -1씩 더하기)한 뒤 해당 set에서 그 노드들을 제거하는 등의 방법으로 이뤄낼 수 있다.

ii) x번 노드에 이미 c 색이 칠해져있는 경우

이 경우 다시 x번 노드에 c를 칠하려는 시도를 할 필요가 없다. 이 경우를 발견하기 위해 위에서 만든 "c 칠하기" 시행 set의 원소 중 x를 서브트리의 원소로 갖는 시행이 있었는지를 찾자. 위와 같이 관리된 set의 각 구간은 서로 겹치는 구간이 없고, set의 특성상 그러한 구간들이 오름차순으로 정렬되어있음을 이용하면 확인해야 할 노드는 많아야 1개이며 그 노드를 빠르게 찾아낼 수 있음을 관찰할 수 있다.

위의 관찰들을 이용해 문제를 해결하자.

아래는 제출한 소스코드이다.

#include <iostream>
#include <vector>
#include <set>
using namespace std;
typedef long long ll;

int N, Q;
vector<int> G[100001];
int dfi[100001];
int childcnt[100001];
int dfsidx = 1;

int dfs(int cur) {
	int ret = 0;
	dfi[cur] = dfsidx++;
	for (auto& nxt : G[cur]) {
		if (dfi[nxt]) continue;
		ret += dfs(nxt);
	}

	childcnt[dfi[cur]] = ret;
	return ret + 1;
}

ll seg[262145];
ll lazy[262145];

void propagate(int L, int R, int sI) {
	seg[sI] += lazy[sI] * (R - L + 1);
	if (L < R) lazy[sI * 2] += lazy[sI], lazy[sI * 2 + 1] += lazy[sI];
	lazy[sI] = 0;
}

ll upd(int L, int R, int qL, int qR, ll qVal, int sI) {
	if (lazy[sI]) propagate(L, R, sI);
	if (qR < L || R < qL) return seg[sI];
	if (qL <= L && R <= qR) {
		lazy[sI] += qVal;
		propagate(L, R, sI);
		return seg[sI];
	}
	return seg[sI] = upd(L, (L + R) / 2, qL, qR, qVal, sI * 2) + upd((L + R) / 2 + 1, R, qL, qR, qVal, sI * 2 + 1);
}

ll qry(int L, int R, int qL, int qR, int sI) {
	if (lazy[sI]) propagate(L, R, sI);
	if (qR < L || R < qL) return 0;
	if (qL <= L && R <= qR) return seg[sI];
	return qry(L, (L + R) / 2, qL, qR, sI * 2) + qry((L + R) / 2 + 1, R, qL, qR, sI * 2 + 1);
}

set<int> st[100001];

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

	cin >> N >> Q;
	for (int i = 1; i < N; i++) {
		int x, y; cin >> x >> y;
		G[x].emplace_back(y);
		G[y].emplace_back(x);
	}

	dfs(1);

	while (Q--) {
		int q; cin >> q;
		if (q == 1) {
			int x, c; cin >> x >> c;
			x = dfi[x];

			if (!st[c].empty()) {
				auto iterT = st[c].lower_bound(x);
				if (iterT == st[c].begin()) {
					if (*iterT == x) continue;
				}
				else {
					iterT = prev(iterT);
					int tmp = *iterT;
					if (tmp <= x && x <= tmp + childcnt[tmp]) continue;
				}
			}

			auto iterL = st[c].lower_bound(x), iterR = st[c].upper_bound(x + childcnt[x]);
			auto iter = iterL;
			while (iter != iterR) {
				int cur = *iter;
				upd(1, N, cur, cur + childcnt[cur], -1, 1);
				iter++;
			}
			st[c].erase(iterL, iterR);
			st[c].insert(x);
			upd(1, N, x, x + childcnt[x], 1, 1);
		}
		else {
			int x; cin >> x;
			x = dfi[x];
			cout << qry(1, N, x, x + childcnt[x], 1) << '\n';
		}
	}
}