1717번 - 집합의 표현(C++)

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

기본적인 유니온 파인드를 사용해서 푸는 문제입니다.

 

첫째줄에 n,m이 주어지고 n개의 원소의 집합이 있으며 m개의 줄만큼 명령이 주어집니다.

0 a b는 union을 1 a b는 두 수가 같은 집합에 속하는지 판별하는것을 의미합니다.

우리는 vector로  "arr[n] = 대표 노드" 이렇게 표현할겁니다.

 

 

초기 상태

모든 노드가 자기 자신을 대표 노드로 가집니다.
즉, arr[i] = i 입니다.

1	2	3	4	5	6	7
1	2	3	4	5	6	7

집합 상태: {1}, {2}, {3}, {4}, {5}, {6}, {7}

1단계: 0 1 3 (union(1,3))

1의 루트 = 1, 3의 루트 = 3 → 3을 1의 대표로 연결.

1	2	3	4	5	6	7
1	2	1	4	5	6	7

집합 상태: {1,3}, {2}, {4}, {5}, {6}, {7}

2단계: 0 3 5 (union(3,5))

3의 루트 = 1, 5의 루트 = 5 → 5를 1에 직접 연결 (경로 압축 적용).

1	2	3	4	5	6	7
1	2	1	4	1	6	7

집합 상태: {1,3,5}, {2}, {4}, {6}, {7}

3단계: 1 1 5 (find(1), find(5))

1의 루트 = 1
5의 루트 = 1 (이미 경로 압축 적용)
따라서 둘 다 최종 루트는 1

결과: YES (같은 집합)

 

코드로 구현해보면

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

// 두 집합을 합치는 함수 (Union 연산)
void unite(vector<int>& arr, int a, int b);

// 특정 원소가 속한 집합의 대표 노드를 찾는 함수 (Find 연산, 경로 압축 포함)
int find(vector<int>& arr,int a);

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

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

    // arr[i] = i의 대표 노드(parent)
    // 초기에는 모든 원소가 자기 자신을 대표로 가짐
    vector<int> arr(n + 1);
    for (int i = 0; i <= n; i++) {
        arr[i] = i;
    }

    // m개의 연산 처리
    for (int i = 0; i < m; i++) {
        int req, a, b;
        cin >> req >> a >> b;

        switch (req)
        {
        case 0: 
            // 0 a b → a와 b가 속한 집합을 합침
            unite(arr, a, b);
            break;

        case 1: 
            // 1 a b → a와 b가 같은 집합인지 확인
            // find로 각각 대표 노드를 찾고 비교
            if (find(arr, a) == find(arr, b)) {
                cout << "YES" << "\n";
            }
            else {
                cout << "NO" << "\n";
            }
            break;

        default:
            break;
        }
    }
    return 0;
}

// unite: 두 원소가 속한 집합을 합치는 함수
// - 먼저 a, b의 대표 노드를 찾음
// - 루트가 다르면 한쪽 루트를 다른 쪽에 붙임
//   (여기서는 단순히 b의 루트를 a의 루트에 연결)
void unite(vector<int>& arr,int a, int b) {
    a = find(arr, a);  // a의 대표 노드
    b = find(arr, b);  // b의 대표 노드
    if (a != b) {
        arr[b] = a;    // b의 루트를 a 밑에 붙임
    }
}

// find: 특정 원소 a의 대표 노드를 찾는 함수
// - 재귀적으로 부모를 따라 올라가면서 루트 탐색
// - 경로 압축(path compression) 기법을 사용
//   → 탐색 과정에서 만나는 모든 노드를 최종 루트로 직접 연결
//   → 이후 탐색이 훨씬 빨라짐
int find(vector<int>& arr,int a) {
    if (arr[a] == a) return a;               // 자기 자신이 루트면 반환
    return arr[a] = find(arr, arr[a]);       // 경로 압축: 루트를 찾으면서 부모를 갱신
}