verify/verify-yosupo-ds/yosupo-weighted-dsu.test.cpp

• View this file on GitHub
• Last update: 2025-03-25 20:32:19+08:00
• Problem: https://judge.yosupo.jp/problem/unionfind_with_potential

Depends on

• src/alfred/config/io-sync-off.hpp
• src/alfred/data_structure/weighted-dsu.hpp
• src/alfred/math/mod-int.hpp

Code

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

#include "../../src/alfred/config/io-sync-off.hpp"
#include "../../src/alfred/data_structure/weighted-dsu.hpp"
#include "../../src/alfred/math/mod-int.hpp"
#include <iostream>

int n, q, opt, u, v, x;
using mint = ModInt<998244353>;

int main(int argc, char const *argv[]) {
    optimizeIO(), std::cin >> n >> q;

    WeightedDSU<mint> dsu(n);

    while (q--) {
        std::cin >> opt >> u >> v;
        if (opt == 0) {
            std::cin >> x;
            std::cout << dsu.merge(u, v, x) << '\n';
        } else if (!dsu.same(u, v)) {
            std::cout << "-1\n";
        } else {
            std::cout << dsu.distance(u, v) << '\n';
        }
    }

    return 0;
}

#line 1 "verify/verify-yosupo-ds/yosupo-weighted-dsu.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/unionfind_with_potential"

#line 1 "src/alfred/config/io-sync-off.hpp"
#include <bits/stdc++.h>

inline void optimizeIO(void) {
    std::ios::sync_with_stdio(false);
    std::cin.tie(NULL), std::cout.tie(NULL);
}
#line 1 "src/alfred/data_structure/weighted-dsu.hpp"



#line 6 "src/alfred/data_structure/weighted-dsu.hpp"

template <class T>
struct WeightedDSU {
    std::vector<int> fa;
    std::vector<T> w;
    WeightedDSU(int n) : fa(n + 1), w(n + 1) {
        std::iota(fa.begin(), fa.end(), 0);
    }
    inline int find(int x) {
        if (fa[x] == x) return x;
        int f = fa[x], f2 = find(f);
        return w[x] += w[f], fa[x] = f2;
    }
    inline bool same(int x, int y) {
        return find(x) == find(y);
    }
    // Given info: a[x] + v = a[y]
    // Returns true if this operation has no conflict, false otherwise.
    inline bool merge(int x, int y, T v) {
        int fx = find(x), fy = find(y);
        if (fx == fy) {
            return w[x] + v == w[y];
        }
        w[fy] = w[x] + v - w[y], fa[fy] = fx;
        return true;
    }
    inline T distance(int x, int y) {
        return find(x), find(y), w[y] - w[x];
    }
};


#line 1 "src/alfred/math/mod-int.hpp"


#ifndef AFMT_MOD_INT
#define AFMT_MOD_INT

#include <iostream>

template <int mod>
inline int down(int x) { return x >= mod ? x - mod : x; }
template <int mod>
inline int up(int x) { return x < 0 ? x + mod : x; }
template <int mod>
struct ModInt {
    int x;
    ModInt(void) = default;
    ModInt(int x) : x(up<mod>(x)) {}
    ModInt(long long x) : x(x % mod) {}
    friend std::istream &operator>>(std::istream &in, ModInt &a) { return in >> a.x; }
    friend std::ostream &operator<<(std::ostream &out, ModInt a) { return out << a.x; }
    friend ModInt operator+(ModInt a, ModInt b) { return down<mod>(a.x + b.x); }
    friend ModInt operator-(ModInt a, ModInt b) { return down<mod>(a.x - b.x + mod); }
    friend ModInt operator*(ModInt a, ModInt b) { return (long long)a.x * b.x; }
    friend ModInt operator/(ModInt a, ModInt b) { return a * ~b; }
    friend ModInt operator^(ModInt a, long long b) {
        ModInt ans = 1;
        for (; b > 0; b >>= 1, a *= a)
            if (b & 1) ans *= a;
        return ans;
    }
    friend ModInt operator~(ModInt a) { return a ^ (mod - 2); }
    friend ModInt operator-(ModInt a) { return down<mod>(mod - a.x); }
    friend ModInt &operator+=(ModInt &a, ModInt b) { return a = a + b; }
    friend ModInt &operator-=(ModInt &a, ModInt b) { return a = a - b; }
    friend ModInt &operator*=(ModInt &a, ModInt b) { return a = a * b; }
    friend ModInt &operator/=(ModInt &a, ModInt b) { return a = a / b; }
    friend ModInt &operator^=(ModInt &a, long long b) { return a = a ^ b; }
    friend ModInt &operator++(ModInt &a) { return a += 1; }
    friend ModInt operator++(ModInt &a, int) {
        ModInt x = a;
        a += 1;
        return x;
    }
    friend ModInt &operator--(ModInt &a) { return a -= 1; }
    friend ModInt operator--(ModInt &a, int) {
        ModInt x = a;
        a -= 1;
        return x;
    }
    friend bool operator==(ModInt a, ModInt b) { return a.x == b.x; }
    friend bool operator!=(ModInt a, ModInt b) { return !(a == b); }
};
template <int mod>
inline void __print(ModInt<mod> x) {
    std::cerr << x;
}

using m998 = ModInt<998244353>;
using m107 = ModInt<1000000007>;

#endif // AFMT_MOD_INT
#line 7 "verify/verify-yosupo-ds/yosupo-weighted-dsu.test.cpp"

int n, q, opt, u, v, x;
using mint = ModInt<998244353>;

int main(int argc, char const *argv[]) {
    optimizeIO(), std::cin >> n >> q;

    WeightedDSU<mint> dsu(n);

    while (q--) {
        std::cin >> opt >> u >> v;
        if (opt == 0) {
            std::cin >> x;
            std::cout << dsu.merge(u, v, x) << '\n';
        } else if (!dsu.same(u, v)) {
            std::cout << "-1\n";
        } else {
            std::cout << dsu.distance(u, v) << '\n';
        }
    }

    return 0;
}

Back to top page