test/library_checker/data_structure/range_parallel_unionfind.test.cpp

• View this file on GitHub
• Last update: 2024-12-06 00:58:38+09:00
• Problem: https://judge.yosupo.jp/problem/range_parallel_unionfind

Depends on

• Range Parallel Union-Find (lib/data_structure/range_parallel_union_find.hpp)
• lib/data_structure/segment_tree_func.hpp
• Union-Find (lib/data_structure/union_find.hpp)
• ModInt (lib/math/modint.hpp)

Code

#define PROBLEM "https://judge.yosupo.jp/problem/range_parallel_unionfind"
#include <iostream>

#include "../../../lib/data_structure/range_parallel_union_find.hpp"
#include "../../../lib/math/modint.hpp"

using namespace std;

using mint = ModInt<998244353>;

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

    int n, q; cin >> n >> q;
    vector<mint> x(n);
    for(int i = 0; i < n; ++i){
        cin >> x[i];
    }

    RangeParallelUnionFind uf(n);
    mint ans = 0;
    while(q--){
        int k, a, b; cin >> k >> a >> b;
        uf.unite(k, a, b, [&](int ra, int rb){
            ans += x[ra] * x[rb];
            x[ra] += x[rb];
        });
        cout << ans << '\n';
    }
}

#line 1 "test/library_checker/data_structure/range_parallel_unionfind.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/range_parallel_unionfind"
#include <iostream>

#line 2 "lib/data_structure/range_parallel_union_find.hpp"

/**
 * @brief Range Parallel Union-Find
 */

#line 2 "lib/data_structure/union_find.hpp"

/**
 * @brief Union-Find
 * @docs docs/data_structure/union_find.md
 */

#include <vector>
#include <cassert>

struct UnionFind{
    int V;
    std::vector<int> par;
    std::vector<int> edg;

    UnionFind(const int N) : V(N), par(N), edg(N){
        for(int i = 0; i < N; ++i){
            par[i] = -1;
            edg[i] = 0;
        }
    }

    int root(int x){
        assert(0 <= x && x < V);
        if(par[x] < 0) return x;
        return par[x] = root(par[x]);
    }

    int unite(int x, int y){
        int rx = root(x);
        int ry = root(y);
        if(rx == ry){
            edg[rx]++;
            return rx;
        }
        if(-par[rx] < -par[ry]) std::swap(rx, ry);
        par[rx] = par[rx] + par[ry];
        par[ry] = rx;
        edg[rx] += edg[ry] + 1;
        return rx;
    }

    bool same(int x, int y){
        int rx = root(x);
        int ry = root(y);
        return rx == ry;
    }

    long long size(int x){
        return -par[root(x)];
    }

    long long edge(int x){
        return edg[root(x)];
    }
};
#line 2 "lib/data_structure/segment_tree_func.hpp"

#line 5 "lib/data_structure/segment_tree_func.hpp"

template <typename T, T (*op)(T, T), T(*e)()>
struct SegTree{
    int _n, n;
    std::vector<T> dat;

    SegTree(int _n) : _n(_n) {
        int x = 1;
        while(_n > x){
            x *= 2;
        }
        n = x;
        dat.resize(n * 2);
        for(int i = 0; i < n * 2; ++i){
            dat[i] = e();
        }
    }
    SegTree(std::vector<T> &v) : _n((int) v.size()) {
        int x = 1;
        while((int) v.size() > x){
            x *= 2;
        }
        n = x;
        dat.resize(n * 2);
        for(int i = 0; i < n; ++i){
            set(i, (i < (int) v.size() ? v[i] : e()));
        }
        build();
    }

private:
    void set(int i, const T &x){ dat[i + n] = x; }

    void build(){
        for(int k = n - 1; k >= 1; k--) dat[k] = op(dat[k * 2], dat[k * 2 + 1]);
    }

public:
    T get(int i) const {
        assert(0 <= i && i < n);
        return dat[i + n];
    }

    void update(int i, const T &x){
        assert(0 <= i && i < n);
        i += n;
        dat[i] = x;
        while(i > 0){
            i >>= 1;  // parent
            dat[i] = op(dat[i * 2], dat[i * 2 + 1]);
        }
    }

    T query(int a, int b){
        assert(0 <= a && a <= b && b <= n);
        T vl = e();
        T vr = e();
        int l = a + n;
        int r = b + n;
        while(l < r){
            if(l & 1) vl = op(vl, dat[l++]);
            if(r & 1) vr = op(dat[--r], vr);
            l >>= 1;
            r >>= 1;
        }
        return op(vl, vr);
    }

    template <class F>
    int max_right(int l, F f) const {
        assert(0 <= l && l <= _n);
        assert(f(e()));
        if(l == _n) return _n;
        l += n;
        T sm = e();
        do{
            while(l % 2 == 0) l >>= 1;
            if(!f(op(sm, dat[l]))){
                while(l < n){
                    l = (2 * l);
                    if(f(op(sm, dat[l]))){
                        sm = op(sm, dat[l]);
                        l++;
                    }
                }
                return l - n;
            }
            sm = op(sm, dat[l]);
            l++;
        }while((l & -l) != l);
        return _n;
    }

    template <class F>
    int min_left(int r, F f) const {
        assert(0 <= r && r <= _n);
        assert(f(e()));
        if(r == 0) return 0;
        r += n;
        T sm = e();
        do{
            r--;
            while(r > 1 && (r % 2)) r >>= 1;
            if(!f(op(dat[r], sm))){
                while(r < n){
                    r = (2 * r + 1);
                    if(f(op(dat[r], sm))){
                        sm = op(dat[r], sm);
                        r--;
                    }
                }
                return r + 1 - n;
            }
            sm = op(dat[r], sm);
        }while((r & -r) != r);
        return 0;
    }
};
#line 9 "lib/data_structure/range_parallel_union_find.hpp"

#line 12 "lib/data_structure/range_parallel_union_find.hpp"

namespace internal_unionfind{
    using i128 = __int128_t;
    constexpr i128 mod = (1LL << 61) - 1;
    constexpr i128 base = 1000003;
    using T = std::pair<i128, i128>;

    T op(T a, T b){
        auto [hash1, len1] = a;
        auto [hash2, len2] = b;
        return std::make_pair((hash1 + hash2 * len1) % mod, (len1 * len2) % mod);
    }

    // {hash, len}
    T e(){
        return std::make_pair(0, 1);
    }
}

struct RangeParallelUnionFind{
    int V;
    std::vector<std::vector<int>> v;
    UnionFind uf;
    SegTree<internal_unionfind::T, internal_unionfind::op, internal_unionfind::e> seg;

private:
    template <typename F>
    void merge(int a, int b, const F &f){
        int ra = uf.root(a), rb = uf.root(b);
        if(ra == rb) return;
        if(uf.unite(ra, rb) == rb) std::swap(ra, rb);
        for(auto e : v[rb]){
            v[ra].push_back(e);
            seg.update(e, std::make_pair(ra, internal_unionfind::base));
        }
        f(ra, rb);
    }

public:
    RangeParallelUnionFind(int n) : V(n), v(n), uf(n), seg(n){
        for(int i = 0; i < n; ++i){
            seg.update(i, std::make_pair(i, internal_unionfind::base));
            v[i].push_back(i);
        }
    }

    // [a, a + k) と [b, b + k) をマージ
    template <typename F>
    void unite(int k, int a, int b, const F &f){
        assert(0 <= a && a + k <= V);
        assert(0 <= b && b + k <= V);
        int prev = 0;
        while(seg.query(a, a + k).first != seg.query(b, b + k).first){
            int ok = k, ng = prev;
            while(abs(ok - ng) > 1){
                int mid = (ok + ng) / 2;
                if(seg.query(a, a + mid).first != seg.query(b, b + mid).first){
                    ok = mid;
                }else{
                    ng = mid;
                }
            }
            merge(a + ok - 1, b + ok - 1, f);
            prev = ok;
        }
    }

    template <typename F>
    void unite(int a, int b, const F &f){
        unite(1, a, b, f);
    }

    int root(int a){
        return uf.root(a);
    }

    bool same(int a, int b){
        return uf.same(a, b);
    }

    long long size(int a){
        return uf.size(a);
    }
};
#line 2 "lib/math/modint.hpp"

#line 5 "lib/math/modint.hpp"

/**
 * @brief ModInt
 * @docs docs/math/modint.md
 */

template <long long Modulus>
struct ModInt{
    long long val;
    static constexpr int mod() { return Modulus; }
    constexpr ModInt(const long long _val = 0) noexcept : val(_val) {
        normalize();
    }
    void normalize(){
        val = (val % Modulus + Modulus) % Modulus;
    }
    inline ModInt &operator+=(const ModInt &rhs) noexcept {
        if(val += rhs.val, val >= Modulus) val -= Modulus;
        return *this;
    }
    inline ModInt &operator-=(const ModInt &rhs) noexcept {
        if(val -= rhs.val, val < 0) val += Modulus;
        return *this;
    }
    inline ModInt &operator*=(const ModInt &rhs) noexcept {
        val = val * rhs.val % Modulus;
        return *this;
    }
    inline ModInt &operator/=(const ModInt &rhs) noexcept {
        val = val * inv(rhs.val).val % Modulus;
        return *this;
    }
    inline ModInt &operator++() noexcept {
        if(++val >= Modulus) val -= Modulus;
        return *this;
    }
    inline ModInt operator++(int) noexcept {
        ModInt t = val;
        if(++val >= Modulus) val -= Modulus;
        return t;
    }
    inline ModInt &operator--() noexcept {
        if(--val < 0) val += Modulus;
        return *this;
    }
    inline ModInt operator--(int) noexcept {
        ModInt t = val;
        if(--val < 0) val += Modulus;
        return t;
    }
    inline ModInt operator-() const noexcept { return (Modulus - val) % Modulus; }
    inline ModInt inv(void) const { return inv(val); }
    ModInt pow(long long n) const {
        assert(0 <= n);
        ModInt x = *this, r = 1;
        while(n){
            if(n & 1) r *= x;
            x *= x;
            n >>= 1;
        }
        return r;
    }
    ModInt inv(const long long n) const {
        long long a = n, b = Modulus, u = 1, v = 0;
        while(b){
            long long t = a / b;
            a -= t * b; std::swap(a, b);
            u -= t * v; std::swap(u, v);
        }
        u %= Modulus;
        if(u < 0) u += Modulus;
        return u;
    }
    friend inline ModInt operator+(const ModInt &lhs, const ModInt &rhs) noexcept { return ModInt(lhs) += rhs; }
    friend inline ModInt operator-(const ModInt &lhs, const ModInt &rhs) noexcept { return ModInt(lhs) -= rhs; }
    friend inline ModInt operator*(const ModInt &lhs, const ModInt &rhs) noexcept { return ModInt(lhs) *= rhs; }
    friend inline ModInt operator/(const ModInt &lhs, const ModInt &rhs) noexcept { return ModInt(lhs) /= rhs; }
    friend inline bool operator==(const ModInt &lhs, const ModInt &rhs) noexcept { return lhs.val == rhs.val; }
    friend inline bool operator!=(const ModInt &lhs, const ModInt &rhs) noexcept { return lhs.val != rhs.val; }
    friend inline std::istream &operator>>(std::istream &is, ModInt &x) noexcept {
        is >> x.val;
        x.normalize();
        return is;
    }
    friend inline std::ostream &operator<<(std::ostream &os, const ModInt &x) noexcept { return os << x.val; }
};
#line 6 "test/library_checker/data_structure/range_parallel_unionfind.test.cpp"

using namespace std;

using mint = ModInt<998244353>;

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

    int n, q; cin >> n >> q;
    vector<mint> x(n);
    for(int i = 0; i < n; ++i){
        cin >> x[i];
    }

    RangeParallelUnionFind uf(n);
    mint ans = 0;
    while(q--){
        int k, a, b; cin >> k >> a >> b;
        uf.unite(k, a, b, [&](int ra, int rb){
            ans += x[ra] * x[rb];
            x[ra] += x[rb];
        });
        cout << ans << '\n';
    }
}

Back to top page