AOJ 1330 Never Wait for Weights

問題

Never Wait for Weights | Aizu Online Judge

N 個の頂点があり, 各頂点には W[i] という特徴量がある。次の二つのクエリが与えられるので適切に処理せよ。

• W[b] - W[a] = w という情報が与えられる。
• a, b を与えるので, 今までの情報で W[b]-W[a] が求められるか調べる。求められる場合はその値を求め, そうでない場合は UNKNOWN と書く。

解法

データ構造をマージする一般的なテクを使います。

W[b]-W[a] の値が分かったとき, a と b がすでに同じグループにいるなら何もしませんが, 同じグループにいない場合は, a と b のグループをマージします。マージするときは,

• b を a のグループにくっつける
• b のグループの重さを a のグループに合わせて更新する

の 2 つの作業をすれば良いです。

int N, M;

int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    while (cin >> N >> M) {
        if (N==0 && M==0) break;
        vector<int> W(N), i2g(N);
        vector<vector<int> > g2i(N);
        // init
        for (int i = 0; i < N; i++) {
            i2g[i] = i;
            g2i[i].push_back(i);
        }
        while (M--) {
            char c;
            cin >> c;
            if (c == '!') {
                int a, b, w;
                cin >> a >> b >> w;
                a--; b--;
                if (i2g[a] == i2g[b]) {
                    assert(W[b]-W[a] == w);
                } else {
                    // merge
                    if (g2i[i2g[a]].size() < g2i[i2g[b]].size()) {
                        swap(a, b);
                        w = -w;
                    }
                    int ga = i2g[a], gb = i2g[b];
                    int diff = (W[a]+w) - W[b];
                    for (int el : g2i[gb]) {
                        i2g[el] = ga; 
                        W[el] += diff;
                    }
                    g2i[ga].insert(g2i[ga].end(), g2i[gb].begin(), g2i[gb].end());
                    g2i[gb].clear();
                }
            } else {
                int a, b;
                cin >> a >> b;
                a--; b--;
                if (i2g[a] != i2g[b]) {
                    cout << "UNKNOWN" << endl;
                } else {
                    cout << W[b]-W[a] << endl;
                }
            }
        }
    }
    return 0;
}