unionfind/test/potentialized_unionfind_modint.test.cpp

• View this file on GitHub
• Last update: 2025-08-05 22:50:49+09:00
• Problem: https://judge.yosupo.jp/problem/unionfind_with_potential

Depends on

• Potentialized UnionFind （重み付き UnionFind） (unionfind/potentialized_unionfind.hpp)

Code

#define PROBLEM "https://judge.yosupo.jp/problem/unionfind_with_potential"

#include "../potentialized_unionfind.hpp"

#include <iostream>
using namespace std;

#include <atcoder/modint>
using mint = atcoder::modint998244353;

int main() {
    cin.tie(nullptr), ios::sync_with_stdio(false);

    int N, Q;
    cin >> N >> Q;
    PotentializedUnionFind<mint> uf(N);

    while (Q--) {
        int t, u, v;
        cin >> t >> u >> v;
        if (t == 0) {
            int x;
            cin >> x;
            if (uf.same(u, v) and uf.diff(v, u) != x) {
                cout << "0\n";
            } else {
                cout << "1\n";
                uf.unite(v, u, x);
            }
        } else {
            if (uf.same(u, v)) {
                cout << uf.diff(v, u).val() << '\n';
            } else {
                cout << "-1\n";
            }
        }
    }
}

#line 1 "unionfind/test/potentialized_unionfind_modint.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/unionfind_with_potential"

#line 2 "unionfind/potentialized_unionfind.hpp"
#include <numeric>
#include <utility>
#include <vector>

// Potentialized UnionFind (Weighted UnionFind)
template <class S> struct PotentializedUnionFind {
    std::vector<int> par, sz;
    std::vector<S> pot;
    PotentializedUnionFind(int N = 0) : par(N), sz(N, 1), pot(N) {
        std::iota(par.begin(), par.end(), 0);
    }
    int find(int x) {
        if (par[x] != x) {
            int r = find(par[x]);
            pot[x] = pot[x] + pot[par[x]], par[x] = r;
        }
        return par[x];
    }
    bool unite(int s, int t, S rel_diff) {
        // Relate s and t by f[t] = f[s] + rel_diff
        // Return false iff contradiction happens.
        rel_diff = rel_diff + weight(s) + (-weight(t));
        if ((s = find(s)) == (t = find(t))) return rel_diff == 0;
        if (sz[s] < sz[t]) std::swap(s, t), rel_diff = -rel_diff;
        par[t] = s, sz[s] += sz[t], pot[t] = rel_diff;
        return true;
    }
    S weight(int x) { return find(x), pot[x]; }
    S diff(int s, int t) { return weight(t) + (-weight(s)); } // return f[t] - f[s]
    int count(int x) { return sz[find(x)]; }
    bool same(int s, int t) { return find(s) == find(t); }
};
#line 4 "unionfind/test/potentialized_unionfind_modint.test.cpp"

#include <iostream>
using namespace std;

#include <atcoder/modint>
using mint = atcoder::modint998244353;

int main() {
    cin.tie(nullptr), ios::sync_with_stdio(false);

    int N, Q;
    cin >> N >> Q;
    PotentializedUnionFind<mint> uf(N);

    while (Q--) {
        int t, u, v;
        cin >> t >> u >> v;
        if (t == 0) {
            int x;
            cin >> x;
            if (uf.same(u, v) and uf.diff(v, u) != x) {
                cout << "0\n";
            } else {
                cout << "1\n";
                uf.unite(v, u, x);
            }
        } else {
            if (uf.same(u, v)) {
                cout << uf.diff(v, u).val() << '\n';
            } else {
                cout << "-1\n";
            }
        }
    }
}

Back to top page