verify/LC-UnionfindWithPotential.test.cpp

View this file on GitHub
Last update: 2025-05-30 20:02:37+09:00
Problem: https://judge.yosupo.jp/problem/unionfind_with_potential

Depends on

Code

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

#include "../Library/modint.hpp"
#include "../Library/DataStructure/WeightedUnionFind.hpp"

int main(){
    cin.tie(0)->sync_with_stdio(false);
    int N, Q; cin >> N >> Q;

    WeightedUnionFind<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)){
                cout << (uf.Diff(u, v) == x) << '\n';
            }
            else{
                cout << uf.Unite(u, v, x) << '\n';
            }
        }
        else{
            if(uf.Same(u, v)){
                cout << uf.Diff(u, v) << '\n';
            }
            else{
                cout << -1 << '\n';
            }
        }
    }
}

#line 1 "verify/LC-UnionfindWithPotential.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/unionfind_with_potential"

#line 2 "Library/modint.hpp"

/**
 * @file modint.hpp
 * @author log K (lX57)
 * @brief modint
 * @version 1.0
 * @date 2023-08-24
 */

#line 2 "Library/Common.hpp"

/**
 * @file Common.hpp
 */

#include <algorithm>
#include <array>
#include <bitset>
#include <cassert>
#include <cstdint>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <limits>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <tuple>
#include <utility>
#include <vector>
using namespace std;

using ll = int64_t;
using ull = uint64_t;

constexpr const ll INF = (1LL << 62) - (3LL << 30) - 1;
#line 12 "Library/modint.hpp"
using namespace std;

const int mod998 = 998244353;
const int mod107 = 1000000007;

template< int mod >
struct ModInt {
    int x;

    ModInt() : x(0) {}

    ModInt(int64_t y) : x(y >= 0 ? y % mod : (mod - (-y) % mod) % mod) {}

    ModInt &operator+=(const ModInt &p) {
        if((x += p.x) >= mod) x -= mod;
        return *this;
    }

    ModInt &operator-=(const ModInt &p) {
        if((x += mod - p.x) >= mod) x -= mod;
        return *this;
    }

    ModInt &operator*=(const ModInt &p) {
        x = (int) (1LL * x * p.x % mod);
        return *this;
    }

    ModInt &operator/=(const ModInt &p) {
        *this *= p.inverse();
        return *this;
    }

    ModInt operator-() const { return ModInt(-x); }

    ModInt operator+(const ModInt &p) const { return ModInt(*this) += p; }

    ModInt operator-(const ModInt &p) const { return ModInt(*this) -= p; }

    ModInt operator*(const ModInt &p) const { return ModInt(*this) *= p; }

    ModInt operator/(const ModInt &p) const { return ModInt(*this) /= p; }

    bool operator==(const ModInt &p) const { return x == p.x; }

    bool operator!=(const ModInt &p) const { return x != p.x; }

    ModInt inverse() const {
        int a = x, b = mod, u = 1, v = 0, t;
        while(b > 0) {
            t = a / b;
            swap(a -= t * b, b);
            swap(u -= t * v, v);
        }
        return ModInt(u);
    }

    ModInt pow(int64_t n) const {
        if(n == 0) return ModInt(1);
        ModInt ret(1), mul(x);
        while(n > 0) {
            if(n & 1) ret *= mul;
            mul *= mul;
            n >>= 1;
        }
        return ret;
    }

    friend ostream &operator<<(ostream &os, const ModInt &p) {
        return os << p.x;
    }

    friend istream &operator>>(istream &is, ModInt &a) {
        int64_t t;
        is >> t;
        a = ModInt< mod >(t);
        return (is);
    }

    static int get_mod() { return mod; }
};

using mint = ModInt<mod998>;
using mint107 = ModInt<mod107>;

using vm = vector<mint>;
using vvm = vector<vector<mint>>;
using vm107 = vector<mint107>;
using vvm107 = vector<vector<mint107>>;
#line 2 "Library/DataStructure/WeightedUnionFind.hpp"

template<typename Abel = int32_t>
class WeightedUnionFind{
    public:
    WeightedUnionFind(int n) : data_(n, -1), weight_(n, Abel{}){}

    int Find(const int k){
        if(data_[k] < 0) return k;
        int r = Find(data_[k]);
        weight_[k] += weight_[data_[k]];
        return data_[k] = r;
    }

    Abel Weight(const int k){
        Find(k);
        return weight_[k];
    }

    Abel Diff(const int x, const int y){
        return Weight(y) - Weight(x);
    }

    bool Same(const int x, const int y){
        return Find(x) == Find(y);
    }

    bool Unite(int x, int y, Abel w){
        w += Weight(x) - Weight(y);
        x = Find(x), y = Find(y);
        if(x == y) return false;
        if(data_[x] > data_[y]) swap(x, y), w = -w;
        data_[x] += data_[y];
        data_[y] = x;
        weight_[y] = w;
        return true;
    }

    private:
    vector<int> data_;
    vector<Abel> weight_;
};
#line 5 "verify/LC-UnionfindWithPotential.test.cpp"

int main(){
    cin.tie(0)->sync_with_stdio(false);
    int N, Q; cin >> N >> Q;

    WeightedUnionFind<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)){
                cout << (uf.Diff(u, v) == x) << '\n';
            }
            else{
                cout << uf.Unite(u, v, x) << '\n';
            }
        }
        else{
            if(uf.Same(u, v)){
                cout << uf.Diff(u, v) << '\n';
            }
            else{
                cout << -1 << '\n';
            }
        }
    }
}