백준 1717번 c++

문제 출처

백준 1717번

문제

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

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

입력

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

출력

1로 시작하는 입력에 대해서 한 줄에 하나씩 YES/NO로 결과를 출력한다. (yes/no 를 출력해도 된다)

테스트 케이스

7 8
0 1 3
1 1 7
0 7 6
1 7 1
0 3 7
0 4 2
0 1 1
1 1 1

answer:
NO
NO
YES

알고리즘 분류

UnionFind

풀이

우선 main 부분입니다. n(집합의 개수 n + 1)과 m(연산의 개수)을 입력받습니다. 0 ~ n이 모두 {n} 형태이므로 반복문을 통해 자기 자신을 가리켜 줍니다. o가 0이면 두집합을 합쳐주고(합집합) o가 1이면 a와 b가 같은 집합에 있는지 확인해 줍니다.

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

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

	vector<int> set(n + 1);
    for (int i = 1; i <= n; i++) 
        set[i] = i;

    for (int i = 0; i < m; i++) {
        int o, a, b;
        cin >> o >> a >> b;
        if (!o)
            unionParent(set, a, b);
        else 
            findParent(set, a, b);
    }
	return 0;
} 

getParent는 원소의 부모를 찾아주는 함수입니다. set[x]에 기록된 숫자가 x와 같다면 자기 자신이 부모이므로 함수를 종료해주고 아니라면 재귀를 통해 부모를 찾아서 반환해 줍니다.

unionFind는 a와 b의 부모를 찾고 기준을 정해(큰 숫자가 부모 or 작은 숫자가 부모) 두 집합을 합쳐줍니다. findParent는 a와 b의 부모를 찾아 비교해준 뒤 같다면 YES를 다르다면 NO를 출력해 줍니다.

int getParent(vector<int>& set, int x) {
    if (set[x] == x) return x;
    return set[x] = getParent(set, set[x]);
}
void unionParent(vector<int>& set, int a, int b) {
    a = getParent(set, a);
    b = getParent(set, b);
    if (a > b) set[a] = b;
    else set[b] = a;
}
void findParent(vector<int>& set, int a, int b) {
    a = getParent(set, a);
    b = getParent(set, b);
    if (a == b) cout << "YES\n";
    else cout << "NO\n";
}

전체 코드

#include<iostream>
#include<vector>

using namespace std;

int getParent(vector<int>& set, int x) {
    if (set[x] == x) return x;
    return set[x] = getParent(set, set[x]);
}
void unionParent(vector<int>& set, int a, int b) {
    a = getParent(set, a);
    b = getParent(set, b);
    if (a > b) set[a] = b;
    else set[b] = a;
}
void findParent(vector<int>& set, int a, int b) {
    a = getParent(set, a);
    b = getParent(set, b);
    if (a == b) cout << "YES\n";
    else cout << "NO\n";
}

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

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

	vector<int> set(n + 1);
    for (int i = 1; i <= n; i++) 
        set[i] = i;

    for (int i = 0; i < m; i++) {
        int o, a, b;
        cin >> o >> a >> b;
        if (!o)
            unionParent(set, a, b);
        else 
            findParent(set, a, b);
    }
	return 0;
}