[백준/BOJ] 백준 17469번 : 트리의 색깔과 쿼리

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

17469번: 트리의 색깔과 쿼리

 

간선을 제거하는 1번 쿼리가 총 N-1번 있으므로, 들어오는 모든 쿼리를 저장해 놓고 거꾸로 모든 정점이 다 끊어진 상황에서 1번 쿼리시 정점의 간선을 연결하는 오프라인 쿼리를 통해 문제를 해결했다.

1번 쿼리시 정점의 연결은 유니온 파인드의 유니온을 통해 정점의 그룹을 합치면서 그룹이 가지고 있는 색깔도 합쳤고, 2번 쿼리시 유니온 파인드의 파인드를 통해 해당 정점 그룹이 가지고 있는 색깔의 개수를 구했다.

 

코드

#include <iostream>
#include <vector>
#include <algorithm>
#include <unordered_set>
using namespace std;

int n, q;
vector<int> result;
vector<pair<int, int>> query;
vector<int> parent_node(100001, 0);
vector<int> parent(100001);
vector<int> rank_size(100001, 1);
unordered_set<int> group_color[100001];
unordered_set<int>::iterator it;

void Pre() {
	for (int i = 0; i < 100001; i++) {
		parent[i] = i;
	}
}

int Find(int u) {

	if (parent[u] == u) {
		return u;
	}

	return parent[u] = Find(parent[u]);
}

void Merge(int u, int v) {

	u = Find(u);
	v = Find(v);

	if (u == v) {
		return;
	}

	//v쪽으로 합치기
	if (rank_size[u] < rank_size[v]) {
		parent[u] = v;

		for (it = group_color[u].begin(); it != group_color[u].end(); it++) {
			group_color[v].insert(*it);
		}
	}

	//u쪽으로 합치기
	else {
		parent[v] = u;

		if (rank_size[u] == rank_size[v]) {
			rank_size[u]++;
		}

		for (it = group_color[v].begin(); it != group_color[v].end(); it++) {
			group_color[u].insert(*it);
		}
	}
}

int main()
{
	cin.tie(NULL);
	ios_base::sync_with_stdio(false);

	Pre();

	cin >> n >> q;

	for (int i = 2; i <= n; i++) {
		int p;
		cin >> p; //i번 정점의 부모 정점

		parent_node[i] = p;
	}

	for (int i = 1; i <= n; i++) {
		int c;
		cin >> c;

		group_color[i].insert(c);
	}

	for (int i = 0; i < n + q - 1; i++) {
		int x, a;

		cin >> x >> a;

		query.push_back(make_pair(x, a));
	}

	//쿼리 순서를 뒤집는다 (간선을 제거하는 1번 쿼리를 총 n-1번 진행하므로, 거꾸로 모든게 다 끊어진 상황에서 1번 쿼리시 간선을 연결하는 방법으로 생각) 
	reverse(query.begin(), query.end());

	for (int i = 0; i < query.size(); i++) {
		int order = query[i].first;
		int node = query[i].second;

		if (order == 1) { //1번 쿼리일때. node와 해당 node의 부모 노드를 유니온
			Merge(node, parent_node[node]);
		}

		else { //2번 쿼리일때
			int node_group = Find(node);
			int this_result = group_color[node_group].size();
			result.push_back(this_result);
		}
	}

	reverse(result.begin(), result.end());

	for (int i = 0; i < result.size(); i++) {
		cout << result[i] << "\n";
	}

	return 0;
}