test/aoj/dsl/dsl_1_b.test.cpp

• View this file on GitHub
• Last update: 2024-11-18 03:54:10+09:00
• Problem: https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=DSL_1_B

Depends on

• Weighted Union-Find (lib/data_structure/weighted_union_find.hpp)

Code

#define PROBLEM "https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=DSL_1_B"
#include <iostream>

#include "../../../lib/data_structure/weighted_union_find.hpp"

using namespace std;

int main(){
    int n, q; cin >> n >> q;
    WeightedUnionFind<int> tree(n);
    while(q--){
        int t; cin >> t;
        if(t == 0){
            int x, y, z; cin >> x >> y >> z;
            tree.unite(x, y, z);
        }else{
            int x, y; cin >> x >> y;
            if(tree.same(x, y)){
                cout << tree.diff(x, y) << "\n";
            }else{
                cout << "?" << "\n";
            }
        }
    }
}

#line 1 "test/aoj/dsl/dsl_1_b.test.cpp"
#define PROBLEM "https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=DSL_1_B"
#include <iostream>

#line 2 "lib/data_structure/weighted_union_find.hpp"

/**
 * @brief Weighted Union-Find
 * @docs docs/data_structure/weighted_union_find.md
 */

#include <vector>
#include <cassert>

template <typename T>
struct WeightedUnionFind{
    int V;
    std::vector<int> par;
    std::vector<T> diff_weight;
    T ex;

    WeightedUnionFind(const int N, const T &e = 0) : V(N), par(N), diff_weight(N), ex(e){
        for(int i = 0; i < N; ++i){
            par[i] = -1;
            diff_weight[i] = e;
        }
    }

    int root(const int x){
        assert(0 <= x && x < V);
        if(par[x] < 0){
            return x;
        }
        int px = par[x];
        int rx = root(par[x]);
        diff_weight[x] = diff_weight[x] + diff_weight[px];
        return par[x] = rx;
    }

    T weight(const int x){
        root(x);
        return diff_weight[x];
    }

    T diff(const int x, const int y){
        return weight(x) - weight(y);
    }

    void unite(const int x, const int y, const T &w){
        int tx = x, ty = y;
        T tw = w;
        int rx = root(x), ry = root(y);
        if(rx == ry) return;

        if(par[rx] < par[ry]){
            std::swap(rx, ry);
            std::swap(tx, ty);
            tw = ex - tw;
        }

        par[ry] = par[rx] + par[ry];
        par[rx] = ry;
        diff_weight[rx] = ex - diff_weight[tx] + tw + diff_weight[ty];
    }

    bool same(const int x, const int y){
        return root(x) == root(y);
    }

    int size(const int x){
        return -par[root(x)];
    }
};
#line 5 "test/aoj/dsl/dsl_1_b.test.cpp"

using namespace std;

int main(){
    int n, q; cin >> n >> q;
    WeightedUnionFind<int> tree(n);
    while(q--){
        int t; cin >> t;
        if(t == 0){
            int x, y, z; cin >> x >> y >> z;
            tree.unite(x, y, z);
        }else{
            int x, y; cin >> x >> y;
            if(tree.same(x, y)){
                cout << tree.diff(x, y) << "\n";
            }else{
                cout << "?" << "\n";
            }
        }
    }
}

Back to top page