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

문제 이해

- 초기에 n+1개의 집합이 있다.

- {0}, {1}, ..., {n}

- 합집합 연산과 두 집합이 같은 집합에 포함되어 있는지 확인하는 연산이 존재한다.

- 집합을 표현하는 프로그램을 작성하라.

 

 

<조건>

- 시간 제한: 2초

- 메모리 제한: 128MB

 

<입력>

- n: 집합의 개수 관련 변수 (1 ~ 1,000,000, 10^6)

- m: 연산의 개수 (1 ~ 100,000, 10^5)

- 연산: 합집합(0 a b), 같은 집합에 포함되어 있는지 확인(1 a b)

 

<출력>

- 1로 시작하는 입력에 대해서 a와 b가 같은 집합에 포함되어 있으면 yes, 있지않으면 no를 출력하라.

 

 

 

문제 풀이

두 원소가 같은 집합에 속하는지 아닌지를 판단하는 유니온 파인드의 기본 문제이다.

 

 

유니온 파인드 알고리즘은 같은 집합에 속하는 원소들을 트리로 구현한다.

 

 

한 집합에 속하는 원소들은 루트 노드가 모두 동일하도록 트리로 구성하고,

루트 노드의 부모는 자기 자신으로 지정한다.

for(int i = 0; i <= n; i++){
    parent[i] = i;
}

 

 

 

이때, 한 집합의 부모 노드를 루트 노드 1개로 통일할경우

제일 마지막 노드에서 루트 노드까지 탐색하는 최악의 경우(O(N))를 면할 수 있다.

 

 

 

두 원소가 같은 집합에 속하는지 확인하는 연산은

각 원소의 루트 노드가 같은지를 체크함으로서 수행된다.

int find(int a){
    if(parent[a] == a) return a;
    return parent[a] = find(parent[a]);
}

 

return 할 때 값을 parent[a]에 대입해줌으로써

각 노드의 부모를 찾으며 모든 노드의 부모를 루트 노드로 수정해주는 것이다.

 

 

 

합연산의 경우 두 원소가 다른 집합에 속할경우,

한 집합의 루트 노드의 부모를 다른 집합의 루트 노드로 변경함으로써 합연산을 수행한다.

void uni(int a, int b){
    int rootA = find(a);
    int rootB = find(b);
    if(rootA == rootB) return;
    parent[rootA] = rootB;
}

 

 

 

전체 코드

#include <iostream>
#include <algorithm>
#include <vector>
#include <queue>
#include <cstring>
#include <cmath>
using namespace std;
#define INF 987654321

int parent [1010101];

int find(int a){
    if(parent[a] == a) return a;
    return parent[a] = find(parent[a]);
}

void uni(int a, int b){
    int rootA = find(a);
    int rootB = find(b);
    if(rootA == rootB) return;
    parent[rootA] = rootB;
}

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

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

    for(int i = 0; i <= n; i++){
        parent[i] = i;
    }

    while(m--){
        int op, a, b;
        cin>>op>>a>>b;
        if(op == 0){
            uni(a, b);
        }else{
            if(find(a) == find(b)){
                cout<<"yes\n";
            }else{
                cout<<"no\n";
            }
        }
    }
    
    return 0;
}

 

 

 

배운 점

유니온 파인드 알고리즘의 기본 코드에 대해 배울 수 있었다!

응용 문제를 풀며 더 익혀보도록 하자.