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

www.acmicpc.net/problem/3830

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

 Union Find를 응용해서 푸는 문제다. !로 시작하는 두 물체간의 무게 비교 결과가 주어질때는 union을 해주고 'diff[a] = a의 무게 - root의 무게'로 정의해서 계산해준다. 두 물체간 비교를 묻는 질의가 들어올때는 두 정점이 서로 Union이 아니면 'UNKNOWN', Union이면 diff[b] - diff[a]의 결과를 출력해주면 된다.

#include <stdio.h>

int N, M;
long long ans;
long long parents[100010];
long long diff[100010];

long long updateParents(long long a)
{
    if (parents[a] == a)
    {
        return a;
    }
    else
    {
        int t = parents[a];
        parents[a] = updateParents(parents[a]);
        diff[a] += diff[t];
        return parents[a];
    }
}

void unionSample(long long a, long long b, long long w)
{
    updateParents(a);
    updateParents(b);
    long long x = diff[b];
    long long y = diff[a];
    a = updateParents(a);
    b = updateParents(b);
    parents[b] = a;
    diff[b] = w + y - x;
}

int main()
{
    while (true)
    {
        scanf("%d%d", &N, &M);
        if (N == 0 && M == 0)
        {
            break;
        }
        for (int i = 0; i <= N; i++)
        {
            parents[i] = i;
            diff[i] = 0;
        }

        for (int i = 0; i < M; i++)
        {
            char q[2];
            scanf("%s", q);
            if (q[0] == '!')
            {
                int a, b, w;
                scanf("%d%d%d", &a, &b, &w);
                unionSample(a, b, w);
            }
            else
            {
                int a, b;
                scanf("%d%d", &a, &b);
                if (updateParents(a) == updateParents(b))
                {
                    printf("%lld\n", diff[b] - diff[a]);
                }
                else
                {
                    printf("UNKNOWN\n");
                }
            }
        }
    }
    return 0;
}