Algorithm in A..Z - Union Find

개념

서로 같은 집합에 속해있는지 확인하고 같은 집합으로 합치거나 집합의 포함된 원소의 수를 알아내는 알고리즘이다.

 

union => 두 원소 a, b가 있을 때 a가 속한 집합과 b가 속한 집합을 합치는 연산

find => 원소 a가 있을 때 a가 속한 집합을 찾는 연산

---

작동원리

초기화

1. 배열을 -1로 초기화한다. (음수의 절대값이 해당 집합에 속안 원소의 갯수가 된다.)

 

find

1. 재귀호출을 통해 배열에 저장된 값이 음수일 때 까지 찾는다 (배열에 저장된 값이 양수면 해당 원소의 부모 원소이고, 음수면 해당 집합의 최고 위치에 있는 인덱스이다.)

2. 배열에 저장된 값이 음수이면 해당 인덱스를 반환한다.

 

union

1. 두 원소의 속한 집합을 find연산으로 찾는다.

2. 두 원소가 속한 집합이 서로 다를 경우 집합 하나에 다른 집합을 더한다 (음수끼리 더해서 집합의 속한 원소의 갯수를 합친다.

3. 더한 집합의 값을 집합의 값으로 바꾼다.

---

시간 복잡도

O(1)

---

문제

 

16562 친구비

www.acmicpc.net/problem/16562

16562번: 친구비

---

코드

#include <bits/stdc++.h>
using namespace std;

vector<int> disjoint;
int findSet(int a) {
    if (disjoint[a] < 0) {
        return a;
    }

    return disjoint[a] = findSet(disjoint[a]);
}

void unionSet(int a, int b) {
    int aa = findSet(a);
    int bb = findSet(b);

    disjoint[aa] += disjoint[bb];
    disjoint[bb] = aa;
}

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

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

    priority_queue<pair<int, int>, vector<pair<int, int>>, greater<>> que;
    for (int i = 1;i <= n;++i) {
        int cost;
        cin >> cost;

        que.emplace(cost, i);
    }

    disjoint.resize(n + 1, -1);
    while (m--) {
        int a, b;
        cin >> a >> b;

        if (findSet(a) != findSet(b)) {
            unionSet(a, b);
        }
    }

    int sum = 0;
    while (!que.empty()) {
        auto now = que.top();que.pop();

        if (findSet(0) != findSet(now.second)) {
            sum += now.first;
            unionSet(0, now.second);
        }
    }

    if (sum > k) {
        cout << "Oh no" << "\n";
    } else {
        cout << sum << "\n";
    }
}