[SWEA] 1849 영준이의 무게측정

출처: SWEA

Approach

Disjoint-Set & Union-Find 

Disjoint-Set & Union-Find

---

#include <stdio.h>

const int MAX_N = 1e6 + 2;
char cmd;
int N, M, TC, a, b, w;
int parent[MAX_N], rank[MAX_N], W[MAX_N];

int Find(int n)
{
    if (parent[n] == n)
        return n;
    W[n] += W[parent[n]];
    return parent[n] = Find(parent[n]);
}

void Union(int a, int b, int w)
{
    int pa = Find(a), pb = Find(b);
    if (pa == pb)
        return;

    if (rank[pa] < rank[pb]) 
    {
        parent[pa] = pb;
        W[pa] = W[b] - W[a] + w;
    }
    else 
    {
        parent[pb] = pa;
        W[pb] = W[a] - W[b] - w;

        if (rank[pa] == rank[pb])
        {
            rank[pa]++;
        }
    }
}

int main(void)
{
    // freopen("input.txt", "r", stdin);
    scanf("%d", &TC);
    for (int tc = 1; tc <= TC;++tc) {
        scanf("%d %d", &N, &M);
        for (int i = 1; i <= N; i++) 
        {
            parent[i] = i;
            W[i] = rank[i] = 0;
        }

        printf("#%d ", tc);
        for (int i = 1; i <= M; i++) {
            scanf(" %c %d %d", &cmd, &a, &b);
            if (cmd == '!') {
                scanf("%d", &w);
                Union(a, b, w);
            }
            else {
                if (Find(a) == Find(b))
                    printf("%d ", W[a] - W[b]);
                else
                    printf("UNKNOWN ");
            }
        }
        printf("\n");
    }
    return 0;
}