AOJ 1330 Never Wait for Weights (ICPC アジア 2012 F) (500 点)

重み付き Union-Find の練習問題！

• 問題へのリンク

問題概要

個の値 があるが、それらの値がわからないので、特定したい。

次の 2 種類のクエリが 個、与えられるので順に処理せよ。

• 3 つの値 が与えられる。 であるという情報が与えられる。
• 2 つの値 を与える。ここまでの情報で の値を特定できるならばその値を出力し、特定できないならば "UNKNOWN" と出力せよ。

各クエリ毎に矛盾が生じることはないことが保証される。なお、1 個目のクエリは ! a b w という形式で入力が与えられ、2 個目のクエリは ? a b という形式で入力が与えられる。

制約

• 
• 

解法

まさに、重み付き Union-Find を verify するためにあるような問題！！

重み付き Union-Find は、Union-Find の各ノードに重みを持たせたデータ構造である。以下のクエリ処理をサポートする。なお、この表では Union-Find のノード x の重みを と書くことにする。

クエリ	処理内容
merge(x, y, w)	となるように x と y を併合する
issame(x, y)	ノード x, y が同じグループにいるかどうかを判定する
diff(x, y)	ノード x, y が同じグループにいるとき、 の値を返す

詳細は次の記事にて。

qiita.com

コード

計算量は 。

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

// 重み付き Union-Find
template<class Abel> struct UnionFind {
    const Abel UNITY_SUM = 0;      // to be set
    vector<int> par;
    vector<Abel> diff_weight;

    UnionFind() { }
    UnionFind(int n) : 
    par(n, -1), diff_weight(n, UNITY_SUM) {}
    
    int root(int x) {
        if (par[x] < 0) return x;
        else {
            int r = root(par[x]);
            diff_weight[x] += diff_weight[par[x]];
            return par[x] = r;
        }
    }
    
    Abel calc_weight(int x) {
        root(x);
        return diff_weight[x];
    }
    
    bool issame(int x, int y) {
        return root(x) == root(y);
    }
    
    bool merge(int x, int y, Abel w = 0) {
        w += calc_weight(x); w -= calc_weight(y);
        x = root(x); y = root(y);
        if (x == y) return false;
        if (par[x] > par[y]) swap(x, y), w = -w; // merge technique
        par[x] += par[y];
        par[y] = x;
        diff_weight[y] = w;
        return true;
    }
    
    Abel diff(int x, int y) {
        return calc_weight(y) - calc_weight(x);
    }
    
    int size(int x) {
        return -par[root(x)];
    }
};

int main() {
    // 入力 ("0 0" が来るまで連続で入力を受け取る)
    int N, M;
    while (cin >> N >> M, N) {
        // 重み付き Union-Find を初期化
        UnionFind<int> uf(N);

        for (int i = 0; i < M; ++i) {
            char c;
            cin >> c;
            if (c == '!') {
                int a, b, w;
                cin >> a >> b >> w;
                --a, --b;

                // w[b] - w[a] = w となるように併合
                uf.merge(a, b, w);
            } else {
                int a, b;
                cin >> a >> b;
                --a, --b;

                if (!uf.issame(a, b)) {
                    // a, b が併合されていない場合は w[b] - w[a] の値は未定
                    cout << "UNKNOWN" << endl;
                } else {
                    cout << uf.diff(a, b) << endl;
                }
            }
        }
    }
}