[백준 1717] 집합의 표현 (C++)

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

문제

---

초기에 n+1개의 집합 {0},{1},{2},…,{n}이 있다. 여기에 합집합 연산과, 두 원소가 같은 집합에 포함되어 있는지를 확인하는 연산을 수행하려고 한다.

집합을 표현하는 프로그램을 작성하시오.

 

입력

---

첫째 줄에 n, m이 주어진다. m은 입력으로 주어지는 연산의 개수이다. 다음 m개의 줄에는 각각의 연산이 주어진다. 합집합은 0 a b의 형태로 입력이 주어진다. 이는 a가 포함되어 있는 집합과, b 가 포함되어 있는 집합을 합친다는 의미이다. 두 원소가 같은 집합에 포함되어 있는지를 확인하는 연산은 1 a b의 형태로 입력이 주어진다. 이는 a와 b가 같은 집합에 포함되어 있는지를 확인하는 연산이다.

 

출력

---

1로 시작하는 입력에 대해서 a와 b가 같은 집합에 포함되어 있으면 "YES" 또는 "yes"를, 그렇지 않다면 "NO" 또는 "no"를 한 줄에 하나씩 출력한다.

제한

---

•  1 ≤ n ≤ 1 000 000
•  1 ≤ m ≤ 100 000
•  0 ≤ a,b ≤ n
•  a, b는 정수
•  a와 b는 같을 수도 있다.

문제 풀이

---

분리집합 문제.

 

 union-find를 구현하면 되는 문제. 입력이 0의 경우는 두 집합을 union()하면 되고, 입력이 1인 경우는 find()를 통해 부모가 같은지 확인하면 된다.

집합의 부모를 저장하는 parents 배열을 생성한다. 처음에는 모든 집합이 자기 자신 하나만 가지고 있으므로 parents[i] = i로 초기화해준다. find 함수의 경우에는 조상의 번호 (=트리에서 루트 노드)를 찾는 함수이다. 부모가 자기 자신 i이면 자기 자신의 번호인 i를 반환해준다. 아니라면 parents 배열을 업데이트하며 조상의 번호를 찾는다. parents 배열은 루트 노드가 아니라 자기 자신의 부모를 가르킨다.

union 함수의 경우에는 두 집합을 합한다. 두 집합의 조상을 find 함수로 찾은 후, 어느 한쪽의 조상의 부모가 다른 한쪽의 조상이 되도록 만든다.

C++의 경우에는 union이 예약어이기 때문에 merge라는 이름을 지어 함수를 만들었다.

 

아래는 코드.

#include <iostream>
using namespace std;
int* parents;

int findParent(int x)
{
	if (parents[x] == x)
	{
		return x;
	}
	else
	{
		return parents[x] = findParent(parents[x]);
	}
}

void merge(int x, int y)
{
	int xParent = findParent(x);
	int yParent = findParent(y);
	if (xParent < yParent)
	{
		parents[yParent] = xParent;
	}
	else
	{
		parents[xParent] = yParent;
	}
}

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

	int N, M;
	int query, a, b;
	cin >> N >> M;
	parents = new int[N+1];
	for (int i = 0; i < N+1; i++)
	{
		parents[i] = i;
	}
	for (int i = 0; i < M; i++)
	{
		cin >> query >> a >> b;
		if (query == 0)
		{
			merge(a, b);
		}
		else
		{
			int aParent = findParent(a);
			int bParent = findParent(b);
			if (aParent == bParent)
			{
				cout << "YES" << "\n";
			}
			else
			{
				cout << "NO" << "\n";
			}
		}
	}
	return 0;
}