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

문제 설명

문제 링크

http://icpc.me/3830

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

풀이

유니온 파인드를 이용하여 풀었다.

대소 관계를 이용하여 무게 차이를 구할 수 있는 것들을 같은 집합으로 묶는다. 이때, dist 배열에 두 개의 무게 차를 저장해 놓는다. dist 배열의 정의는 dist [i]는 i와 i의 부모의 거리이다. 또한 두 집합의 루트 사이의 거리를 정의해줘야 하는데, a, b를 union 할 때 a의 루트를 pa, b의 루트를 pb라고 하면, dist [pb] = (a부터 루트까지 거리) - (b부터 루트까지 거리) + (a와 b의 무게차)로 정의할 수 있다.

a와 b의 무게 차이를 계산할 수 있는지 여부는 a와 b가 같은 집합에 있는지 확인하면 되고, 그 무게 차이는 Root(x)를 x부터 x의 루트까지의 거리라고 한다면, Roor(b) - Root(a)라고 할 수 있다.

위와 같이 문제를 해결하려면 find함수를 실행할 때 경로를 압축하면 안 된다. 그렇기 때문에 큰 집합에 작은 집합을 union 해줘 트리의 높이가 logN보다 작게 유지시켜줘야 한다.

고난과 역경끝에 15트만에 겨우 맞았다..

코드

#include <iostream>
#include <algorithm>

using namespace std;
using ll = long long;

const ll N = 1e5 + 5;
ll par[N], dist[N], lev[N];

ll find(ll x) {
    if (x == par[x]) return x;
    return find(par[x]);
}

ll rootDist(ll x) {
    ll r = find(x);
    ll ret = 0;
    while (x != r) {
        ret += dist[x];
        x = par[x];
    }
    return ret;
}

void merge(ll x, ll y, ll c) {
    ll a = find(x);
    ll b = find(y);
    if (lev[a] < lev[b]) {
        swap(a, b);
        swap(x, y);
        c = -c;
    }
    if (a == b) return;
    dist[b] = rootDist(x) - rootDist(y) + c;
    par[b] = a;
    if (lev[a] == lev[b]) lev[a]++;
}

void solve() {
    ios::sync_with_stdio(0); cin.tie(0);
    ll n, m; cin >> n >> m;
    if (n == 0 && m == 0) exit(0);
    for (ll i = 1; i <= n; i++) {
        par[i] = i;
        dist[i] = 0;
        lev[i] = 1;
    }
    while (m--) {
        char c; cin >> c;
        if (c == '!') {
            ll x, y, c; cin >> x >> y >> c;
            merge(x, y, c);
        } else {
            ll x, y; cin >> x >> y;
            if (find(x) == find(y)) {
                cout << rootDist(y) - rootDist(x) << '\n';
            }
            else cout << "UNKNOWN\n";
        }
    }
}

int main() {
    while (1) solve();    
}