집합의 표현

백준 1717번: 집합의 표현

아이디어

유니온 파인드 (분리 집합)이라고 부르는 자료구조가 있다.
각 원소는 모두 트리의 루트이고, 이를 합치는 연산(union)과 루트를 확인하는 연산(find) 두가지가 존재한다.

원래는 parent에 자기 부모를 저장해서 주루루룩 타고 올라가는 방식이지만, 이렇게 하면 시간복잡도가 스레기라 그냥 루트를 저장한다.
그래서 find 함수를 잘 보면 마지막에 찾은 루트를 자기 parent에 저장한다.

Union 함수(소문자는 예약어임)는 find 함수로 각각 루트를 찾고, 그 루트의 부모를 다른 루트로 한다. 이렇게 하면 합쳐지기 전에 루트였던 원소의 자식들은 find 함수가 실행됐을 때 parent에 저장된 값이 새로운 루트로 바뀌게 된다. 두 집합이 합쳐지는 것이다!

코드

#include <bits/stdc++.h>

using namespace std;

int N, M;

typedef struct _DisjSet {
    int parent[1000001] = {0,};
    
    void init(int n) {
        for (int i = 0; i < n+1; i++) {
            parent[i] = i;
        }
    }
    
    int find(int n) {
        if (parent[n] == n) return n;
        return parent[n] = find(parent[n]);
    }
    
    void Union(int n1, int n2) {
        n1 = find(n1);
        n2 = find(n2);
        if (n1 == n2) return;
        parent[n2] = n1;
    }
    
} DisjSet;

int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cin >> N >> M;
    
    DisjSet d;
    d.init(N);
    
    for (int i = 0; i < M; i++) {
        int f, a, b;
        cin >> f >> a >> b;
        if (f) {
            if (d.find(a) == d.find(b)) {
                cout << "YES\n";
            }
            else {
                cout << "NO\n";
            }
        }
        else {
            d.Union(a, b);
        }
    }
    return 0;
}

여담

나동빈씨 책에서 보고 파이썬으로는 짜봤는데 C++로는 처음 짜봄. 사실 별 다를게 없다.