백준 17469번, 트리의 색깔과 쿼리 풀이

문제

백준 17469번

코드업 3273번

풀이

Union-Find 자료구조를 사용하여 트리를 관리해줍니다.

UF 자료구조 특성상 간선을 제거하는데 상당히 느립니다.

입력을 다 받고 이를 거꾸로 봐주면서 쿼리를 처리하면, 간선을 제거하는 것이 아니라 합치는 연산으로 볼 수 있습니다.

 

이를 합쳐주면서 "색깔"이라는 요소를 같이 관리해야 합니다.

색깔의 개수를 출력할 때는 중복되는 요소가 없어야 하니 set을 사용하면 될 것 같습니다.

 

여기서 백준 13306번 트리와는 다른 테크닉이 하나 더 필요합니다.

 

small to large입니다.

 

Union by rank와 유사하게 간선을 합치는 연산을 처리할 때, set 역시 합치는 연산을 할 것 입니다.

이 set을 작은 것에서 큰쪽으로 옮겨주면 모든 원소가 최대 $ log N $만큼 이동합니다.

 

https://open.kakao.com/o/gPS4w1xc

역순으로 구성된 매 질의에 대해 답한 것을 다시 역순으로 출력해주면 정답입니다.

소스코드

#include <bits/stdc++.h>

using namespace std;

vector<int> UF(100001), parent(100001);
vector<set<int>> color(100001);

int find(int a) {
	return a == UF[a] ? a : UF[a] = find(UF[a]);
}

int main() {
	ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);
	int n, q;
	cin >> n >> q;

	for (int i = 1; i < 100001; i++) UF[i] = i;

	for (int i = 2; i <= n; i++) cin >> parent[i];
	for (int i = 1; i <= n; i++) {
		int a; cin >> a;
		color[i].insert(a);
	}

	stack<pair<int, int>> query;
	for (int i = 1; i < q + n; i++) {
		int a, b; cin >> a >> b;
		query.push({a, b});
	}

	stack<string> ans;
	while(!query.empty() && q) {
		auto a = query.top(); query.pop();
		if (a.first == 1 && q--) { //merge
			int k = find(parent[a.second]);
			if (color[a.second].size() > color[k].size()) color[a.second].swap(color[k]);
			for (int b : color[a.second]) color[k].insert(b);
			UF[a.second] = find(parent[a.second]);
		} else ans.push(to_string(color[find(a.second)].size()));
	}

	string res;
	while (!ans.empty()) {
		res += ans.top() + "\n"; ans.pop();
	}
	cout << res;
}