test/union-find-tree/PotentializedUnionFindTree.test.cpp

• View this file on GitHub
• Last update: 2023-06-17 04:07:32+09:00
• Problem: http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=DSL_1_B

Depends on

• PotentializedUnionFindTree - ポテンシャル付きUnionFind木 (lib/41-union-find-tree/PotentializedUnionFindTree.cpp)

Code

#define PROBLEM "http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=DSL_1_B"

#include <vector>
#include <iostream>
#include <numeric>
using namespace std;
#include "../../lib/41-union-find-tree/PotentializedUnionFindTree.cpp"

int main(){
    cin.tie(0);ios::sync_with_stdio(false);
    int N,Q; cin >> N >> Q;
    PotentializedUnionFindTree<int> uf(N);
    while(Q--) {
        int q; cin >> q;
        if(q) {
            int x,y; cin >> x >> y;
            if(uf.connected(y,x)) cout << uf.diff(x,y) << endl;
            else cout << "?" << endl;
        }
        else {
            int x,y,z; cin >> x >> y >> z;
            uf.merge(x,y,z);
        }
    }
}

#line 1 "test/union-find-tree/PotentializedUnionFindTree.test.cpp"
#define PROBLEM "http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=DSL_1_B"

#include <vector>
#include <iostream>
#include <numeric>
using namespace std;
#line 1 "lib/41-union-find-tree/PotentializedUnionFindTree.cpp"
/*
 * @title PotentializedUnionFindTree - ポテンシャル付きUnionFind木
 * @docs md/union-find-tree/PotentializedUnionFindTree.md
 */
template<class T> class PotentializedUnionFindTree {
    vector<int> parent,rank;
    vector<T> potential;
    inline int root(int n) {
        if (parent[n] == n) {
            return n;
        }
        else {
            int r = root(parent[n]);
            potential[n] += potential[parent[n]];
            return parent[n] = r;
        }
    }
    inline T updated_dist(int x) {
        root(x);
        return potential[x];
    }
public:
    PotentializedUnionFindTree(int N = 1, T ini = 0) : parent(N),rank(N,1),potential(N,ini) {
        iota(parent.begin(),parent.end(),0);
    }
    inline bool connected(int n, int m) {
        return root(n) == root(m);
    }
    /*
    * potential[m] = potential[n] + dとなるようにマージする
    */
    bool merge(int n, int m, T d=0) {
        d += updated_dist(n);
        d -= updated_dist(m);
        n = root(n); m = root(m);
        if (n == m) return false;
        if (rank[n] < rank[m]) {
            swap(n, m);
            d = -d;
        }
        if (rank[n] == rank[m]) ++rank[n];
        parent[m] = n;
        potential[m] = d;
        return true;
    }
    /*
    * potential[m] - potential[n]
    */
    T diff(int n, int m) {
        return updated_dist(m) - updated_dist(n);
    }
    inline int operator[](int n) {
        return root(n);
    }
    inline void print() {
        for(int i = 0; i < parent.size(); ++i) cout << root(i) << " ";
        cout << endl;
    }
};
#line 8 "test/union-find-tree/PotentializedUnionFindTree.test.cpp"

int main(){
    cin.tie(0);ios::sync_with_stdio(false);
    int N,Q; cin >> N >> Q;
    PotentializedUnionFindTree<int> uf(N);
    while(Q--) {
        int q; cin >> q;
        if(q) {
            int x,y; cin >> x >> y;
            if(uf.connected(y,x)) cout << uf.diff(x,y) << endl;
            else cout << "?" << endl;
        }
        else {
            int x,y,z; cin >> x >> y >> z;
            uf.merge(x,y,z);
        }
    }
}

Back to top page