AOJ 1330 Never Wait for Weights

解法

基本は union-find をいじるだけ.
各ノードは,親からの相対重さを持つ.
クエリで a, b, w が与えられたとする. a と b の代表元が異なる場合は,それぞれの集合の根を r1, r2,各ノードの真の重みを W とおくと,

  • W[r1] - W[r2] = -w + (b の r2 からの相対重み) - (a の r1 からの相対重み)

であることが(簡単な計算で)わかる.
あとはこの処理を union-find 木の処理に追加すればよい.

ソースコード

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

constexpr int INF = 1e9;

class union_find {
public:
    union_find(int N)
        : par(N, -1), rank(N), w(N)
    {}

    pair<int, int> root(int x) {
        if(par[x] < 0) {
            return make_pair(x, 0);
        }
        pair<int, int> p = root(par[x]);
        w[x] += p.second;
        p.second = w[x];
        par[x] = p.first;
        return p;
    }

    void unite(int x, int y, int v) {
        int rx = root(x).first, ry = root(y).first;
        if(rx == ry) {
            return;
        }
        if(rank[rx] > rank[ry]) {
            par[ry] = rx;
            w[ry] = v - w[y] + w[x];
        } else {
            par[rx] = ry;
            w[rx] = -v + w[y] - w[x];
            if(rank[rx] == rank[ry]) {
                rank[ry]++;
            }
        }
    }

    int calc(int x, int y) {
        if(root(x).first != root(y).first) {
            return INF;
        }
        return w[y] - w[x];
    }
        
private:
    vector<int> par;
    vector<int> rank;
    vector<int> w;
};

int main() {
    int N, M;
    while(cin >> N >> M, N) {
        union_find uf(N);
        for(int i=0; i<M; ++i) {
            char op; cin >> op;
            if(op == '!') {
                int a, b, v;
                cin >> a >> b >> v;
                a--; b--;
                uf.unite(a, b, v);
            } else {
                int a, b;
                cin >> a >> b;
                a--; b--;
                if(uf.calc(a, b) == INF) {
                    cout << "UNKNOWN" << endl;
                } else {
                    cout << uf.calc(a, b) << endl;
                }
            }
        }
    }
}