취미로 코딩하기

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

이번에 볼 문제는 백준 18116번 문제인 로봇 조립이다.
문제는 아래 링크를 확인하자.

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

하나의 부품은 하나의 로봇에만 들어가므로, disjoint set 자료구조를 이용하면 문제를 해결할 수 있다.

각 쿼리마다 union과 find 연산을 잘 활용해주고, union 연산을 할 때 해당 부품이 들어가는 로봇에 들어가는 알고 있는 부품의 수를 잘 갱신해주자.

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

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

int uf[1000001];
int cnt[1000001];
int depth[1000001];

int findf(int x) {
	if (x == uf[x]) return x;
	else return uf[x] = findf(uf[x]);
}

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

	for (int i = 1; i <= 1000000; i++) uf[i] = i;
	for (int i = 1; i <= 1000000; i++) cnt[i] = 1;
	int Q; cin >> Q;
	while (Q--) {
		char c; cin >> c;
		if (c == 'I') {
			int x, y; cin >> x >> y;
			x = findf(x), y = findf(y);
			if (x == y) continue;
			if (depth[x] < depth[y]) swap(x, y);
			if (depth[x] == depth[y]) depth[x]++;
			cnt[x] += cnt[y];
			uf[y] = x;
		}
		else {
			int x; cin >> x;
			cout << cnt[findf(x)] << '\n';
		}
	}
}