백준 1717 - 집합의 표현 (feat. 분리 집합)

문제링크

• 백준 1717

문제설명

• 초기에 n+1개의 집합 {0}, {1}, {2}, … , {n}이 있다.
• 여기에 합집합 연산과, 두 원소가 같은 집합에 포함되어 있는지를 확인하는 연산을 수행하려고 한다.
• 집합을 표현하는 프로그램을 작성하시오.
• 첫째 줄에 n, m이 주어진다.
• m은 입력으로 주어지는 연산의 개수이다. 다음 m개의 줄에는 각각의 연산이 주어진다. 합집합은 0 a b의 형태로 입력이 주어진다. 이는 a가 포함되어 있는 집합과, b가 포함되어 있는 집합을 합친다는 의미이다.
• 두 원소가 같은 집합에 포함되어 있는지를 확인하는 연산은 1 a b의 형태로 입력이 주어진다. 이는 a와 b가 같은 집합에 포함되어 있는지를 확인하는 연산이다.

문제접근

• 분리 집합 (Disjoint Set) : Union-Find
• 이 문제는 루트 노드부터 자식 노드까지의 경로에는 관심이 없으므로 모든 자식 노드들의 부모 노드를 루트 노드로 설정한다. -> 경로 단축

구현

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

int FindRoot(vector<int>& parent, int cnode)
{
	if (parent[cnode] == cnode)
	{
		return cnode;
	}
	else
	{
		int rNode = FindRoot(parent, parent[cnode]);
		parent[cnode] = rNode;
		return rNode;
	}
}

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

	int n, m;
	cin >> n >> m;

	int op, a, b;
	vector<int> parent(n + 1);
	for (int i = 0; i <= n; i++)
	{
		parent[i] = i;
	}

	for (int i = 0; i < m; i++)
	{
		cin >> op >> a >> b;
		int aRoot = FindRoot(parent, a);
		int bRoot = FindRoot(parent, b);
		if (op == 0)
		{
			if (aRoot != bRoot) // 두 수의 그룹이 서로 다른 경우
			{
				if (aRoot < bRoot)
				{
					parent[bRoot] = aRoot;
					FindRoot(parent, b); // update parent node to root node
				}
				else
				{
					parent[aRoot] = bRoot;
					FindRoot(parent, a); // update parent node to root node
				}
			}
		}
		else if (op == 1)
		{
			if (aRoot == bRoot)
			{
				cout << "YES\n";
			}
			else
			{
				cout << "NO\n";
			}
		}
	}
}