test/weighted_unionfind.test.cpp

• View this file on GitHub
• Last update: 2022-08-29 11:38:54+09:00
• Problem: https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=1330

Depends on

• data_structures/weighted_unionfind.cpp

Code

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

#define PROBLEM "https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=1330"
#include "../data_structures/weighted_unionfind.cpp"

void solve(int n, int m) {
    WeightedUnionFind<int, 0> wuf(n);
    int a, b, w;
    for (int i = 0; i < m; i++) {
        char q;
        cin >> q;
        if (q == '!') {
            cin >> a >> b >> w;
            wuf.unite(--a, --b, w);
        } else {
            cin >> a >> b;
            if (!wuf.same(--a, --b))
                cout << "UNKNOWN\n";
            else
                cout << wuf.diff(a, b) << '\n';
        }
    }
}

int main() {
    cin.tie(nullptr);
    ios::sync_with_stdio(false);
    cout << fixed << setprecision(20);

    int n, m;
    while (1) {
        cin >> n >> m;
        if (n == 0 && m == 0) break;
        solve(n, m);
    }

    return 0;
}

#line 1 "test/weighted_unionfind.test.cpp"
#include <bits/stdc++.h>
using namespace std;

#define PROBLEM "https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=1330"
#line 3 "data_structures/weighted_unionfind.cpp"
using namespace std;

// WeightedUnionFind
// 通常のUnionFindに加えて, 同じグループにいるノードの親に対する重みも
// 管理するデータ構造.
template <typename Abel, Abel unit>
struct WeightedUnionFind {
    int n;
    vector<int> parents;
    vector<Abel> diff_weight;

    WeightedUnionFind(int n) : n(n) {
        parents.assign(n, -1);
        diff_weight.assign(n, unit);
    }

    // find
    // xの親を返す. 経路圧縮するため, diff_weightも更新する
    // 制約: 0 <= x < n
    int find(int x) {
        if (parents[x] < 0)
            return x;
        else {
            int p = find(parents[x]);
            diff_weight[x] += diff_weight[parents[x]];
            return parents[x] = p;
        }
    }

    // weight
    // xのparentからの重みを返す
    // 制約: 0 <= x < n
    Abel weight(int x) {
        find(x);
        return diff_weight[x];
    }

    // diff
    // xとyが同じグループにいる時, xに対するyの重みを返す.
    // xとyが同じグループにいない時, 返す値には意味がない.
    // 制約: 0 <= x,y < n, xとyは同じグループ
    Abel diff(int x, int y) {
        return weight(y) - weight(x);
    }

    // unite
    // weight(y) = weight(x) + wとなるようにxとyを併合する.
    // もしすでにxとyが同じグループでweight(y) != weight(x) + wであればfalseを返す.
    // そうでない場合はtrueを返す.
    // 制約: 0 <= x,y < n
    bool unite(int x, int y, Abel w) {
        w += weight(x);
        w -= weight(y);

        x = find(x);
        y = find(y);
        if (x == y) {
            if (diff(x, y) == w)
                return true;
            else
                return false;
        }

        if (parents[x] > parents[y]) {
            swap(x, y);
            w *= -1;
        }

        parents[x] += parents[y];
        parents[y]     = x;
        diff_weight[y] = w;
        return true;
    }

    // same
    // xとyが同じかどうか判定する
    // 制約: 0 <= x,y < n
    bool same(int x, int y) {
        return find(x) == find(y);
    }

    // size
    // xと同じグループのサイズを返す.
    // 制約: 0 <= x < n
    int size(int x) {
        return -parents[find(x)];
    }
};
#line 6 "test/weighted_unionfind.test.cpp"

void solve(int n, int m) {
    WeightedUnionFind<int, 0> wuf(n);
    int a, b, w;
    for (int i = 0; i < m; i++) {
        char q;
        cin >> q;
        if (q == '!') {
            cin >> a >> b >> w;
            wuf.unite(--a, --b, w);
        } else {
            cin >> a >> b;
            if (!wuf.same(--a, --b))
                cout << "UNKNOWN\n";
            else
                cout << wuf.diff(a, b) << '\n';
        }
    }
}

int main() {
    cin.tie(nullptr);
    ios::sync_with_stdio(false);
    cout << fixed << setprecision(20);

    int n, m;
    while (1) {
        cin >> n >> m;
        if (n == 0 && m == 0) break;
        solve(n, m);
    }

    return 0;
}

Back to top page