Alfred's CP Library

This documentation is automatically generated by competitive-verifier/competitive-verifier

View on GitHub

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

Depends on

Code

// competitive-verifier: 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"
// competitive-verifier: 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"



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

namespace ModIntImpl {
    template <int mod>
    inline int up(int x) { return x < 0 ? x + mod : x; }
    template <int mod>
    inline int down(int x) { return x >= mod ? 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) { return a.pow(b); }
        friend ModInt operator^(ModInt a, int b) { return a.pow(b); }
        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.x != b.x; }
        inline ModInt pow(long long b) noexcept {
            ModInt ans = 1, a = x;
            for (b %= (mod - 1); b > 0; b >>= 1, a *= a)
                if (b & 1) ans *= a;
            return ans;
        }
        inline ModInt pow(int b) noexcept {
            ModInt ans = 1, a = x;
            for (; b > 0; b >>= 1, a *= a)
                if (b & 1) ans *= a;
            return ans;
        }
    };
}

using ModIntImpl::ModInt;

template <int mod>
inline void __print(ModInt<mod> x) {
    std::cerr << x;
}

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


#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;
}

Test cases

Env Name Status Elapsed Memory
g++ example_00 :heavy_check_mark: AC 5 ms 4 MB
g++ max_random_00 :heavy_check_mark: AC 56 ms 4 MB
g++ max_random_01 :heavy_check_mark: AC 55 ms 5 MB
g++ max_random_02 :heavy_check_mark: AC 60 ms 4 MB
g++ path_00 :heavy_check_mark: AC 56 ms 8 MB
g++ path_01 :heavy_check_mark: AC 50 ms 5 MB
g++ path_02 :heavy_check_mark: AC 62 ms 9 MB
g++ path_03 :heavy_check_mark: AC 58 ms 5 MB
g++ random_00 :heavy_check_mark: AC 40 ms 4 MB
g++ random_01 :heavy_check_mark: AC 40 ms 5 MB
g++ random_02 :heavy_check_mark: AC 37 ms 4 MB
g++ random_03 :heavy_check_mark: AC 11 ms 5 MB
g++ random_04 :heavy_check_mark: AC 31 ms 4 MB
g++ random_05 :heavy_check_mark: AC 42 ms 4 MB
g++ random_06 :heavy_check_mark: AC 30 ms 5 MB
g++ random_07 :heavy_check_mark: AC 8 ms 4 MB
g++ random_08 :heavy_check_mark: AC 16 ms 4 MB
g++ random_09 :heavy_check_mark: AC 55 ms 4 MB
Back to top page