[BOJ 3830] 교수님은 기다리지 않는다

문제

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

3830번: 교수님은 기다리지 않는다

 

---

알고리즘

• Union-Find

---

풀이

배열 2개를 만들어 다음과 같이 정의합니다:

 

• \(uf[i]:\) \(i\)가 속한 집합의 루트 노드
• \(dist[i]:\) \(i - i\)가 속한 집합의 루트 노드

UNKNOWN인지 아닌지는 \(find\) 함수를 사용해 같은 집합에 있는지 아닌지 확인하면 쉽게 알 수 있습니다.

 

입력 받은 \(a, b\)의 차를 구할 때는 \(dist[b] - dist[a]\)를 구하면 됩니다.

여기서 \(dist[x]\)는 \(find\) 함수를 응용해서 \(x\)의 루트 노드를 구할 때마다 갱신해 주면 됩니다.

현 \(x\)에 대하여 \(dist[x]\)에는 \(x\)와 \(uf[x]\)의 차가 저장되어 있으므로 거기에 \(uf[x]\)와 \(uf[uf[x]]\)의 차, 즉 \(dist[uf[x]]\)를 더해주면 됩니다. 그럼 \(x\)의 최종적인 루트 노드에 대한 정보를 얻을 수 있습니다.

int find(int x) {
    if (uf[x] < 0) return x;
    int root = find(uf[x]);
    dist[x] += dist[uf[x]];
    return uf[x] = root;
}

두 집합을 입력으로 인해 합칠 때도 있는데 이때는 \(merge\) 함수를 사용합니다.

\(merge\)하기 전 \(a\)와 \(b\)의 루트 노드를 \(r_a, r_b\)라고 정의하겠습니다.

하나로 합쳐야 하므로 \(uf[r_b]\)에 \(r_a\)를 저장하고(\(r_b\)의 루트는 \(r_a\)가 됩니다.), \(\ dist[r_b]\)에는 \(dist[a] + w\), 즉 \(b\)와 \(r_a\)의 차에서 \(b\)와 \(r_b\)의 차인 \(dist[b]\)를 빼주면 \(r_b\)와 \(r_a\)의 차를 구할 수 있으므로 그 값을 저장합니다. (\(a\)와 \(b\)는 바뀌어도 상관없습니다.)

void merge(int a, int b, int w) {
    int root_a = find(a);
    int root_b = find(b);
    if (root_a == root_b) return;
    dist[root_b] = dist[a] - dist[b] + w;
    uf[root_b] = root_a;
}

여기서 \(r_a\)의 집합에 있었던 무게와 \(r_b\)의 집합에 있어던 무게의 차에 대해서는 들어온 쿼리에 대해 UNKNOWN인지 아닌지 확인하기 위해 \(find\) 함수를 호출하므로, 자연스럽게 갱신이 가능합니다.

 

---

전체 코드

#include <bits/stdc++.h>
#define x first
#define y second
#define pb push_back
#define all(x) (x).begin(), (x).end()
#define rall(x) (x).rbegin(), (x).rend()
using namespace std;

typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;

const int MXN = 100001;
int uf[MXN];
ll dist[MXN];

int find(int x) {
    if (uf[x] < 0) return x;
    int root = find(uf[x]);
    dist[x] += dist[uf[x]];
    return uf[x] = root;
}

void merge(int a, int b, int w) {
    int root_a = find(a);
    int root_b = find(b);
    if (root_a == root_b) return;
    dist[root_b] = dist[a] - dist[b] + w;
    uf[root_b] = root_a;
}

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

    while (1) {
        memset(uf, -1, sizeof(uf));
        memset(dist, 0, sizeof(dist));
        int N, M;
        cin >> N >> M;
        if (!N && !M) break;
        while (M--) {
            char c;
            int a, b, w;
            cin >> c;
            if (c == '!') {
                cin >> a >> b >> w;
                merge(a, b, w);
            }
            else {
                cin >> a >> b;
                if (find(a) != find(b)) cout << "UNKNOWN\n";
                else cout << dist[b] - dist[a] << '\n';
            }
        }
    }

    return 0;
}