library/datastructure/unionfind/PotentialUnionFind.hpp

View this file on GitHub
Last update: 2025-11-24 18:49:07+09:00
Include: #include "library/datastructure/unionfind/PotentialUnionFind.hpp"

Depends on

build/pch/stdc++.hpp

Verified with

Code

#pragma once
template <typename AbelGroup> class PotentialUnionFind {
    using T = typename AbelGroup::value_type;
    int n, num;
    std::vector<int> sz, parent;
    std::vector<T> potential; // parent[x] を基準とした時の x の値
  public:
    PotentialUnionFind() = default;
    PotentialUnionFind(int n)
        : n(n), num(n), sz(n, 1), parent(n, 0),
          potential(n, AbelGroup::unit()) {
        assert(AbelGroup::commute);
        std::iota(parent.begin(), parent.end(), 0);
    }

    std::pair<int, T> from_root(int x) {
        if (x == parent[x])
            return {x, AbelGroup::unit()};
        auto [r, add] = from_root(parent[x]);
        parent[x] = r;
        AbelGroup::Rchop(potential[x], add);
        return {r, potential[x]};
    }

    int leader(int x) { return from_root(x).first; }

    bool same(int x, int y) {
        assert(0 <= x and x < n and 0 <= y and y < n);
        return leader(x) == leader(y);
    }

    bool merge(int x, int y, T d) {
        // potential[y]-potential[x]=d にする
        // 矛盾する場合は変更はせず false を返す
        assert(0 <= x and x < n and 0 <= y and y < n);
        auto [rx, dx] = from_root(x);
        auto [ry, dy] = from_root(y);
        AbelGroup::Rchop(d, dx);
        AbelGroup::Rchop(d, AbelGroup::inverse(dy));
        if (rx == ry)
            return d == AbelGroup::unit();
        if (sz[rx] < sz[ry]) {
            std::swap(rx, ry);
            d = AbelGroup::inverse(d);
        }
        sz[rx] += sz[ry];
        parent[ry] = rx;
        potential[ry] = d;
        num--;
        return true;
    }

    std::optional<T> diff(int x, int y) {
        // x を基準とする
        auto [rx, dx] = from_root(x);
        auto [ry, dy] = from_root(y);
        if (rx != ry)
            return std::nullopt;
        return AbelGroup::op(dy, AbelGroup::inverse(dx));
    }

    int size(const int x) {
        assert(0 <= x and x < n);
        return sz[leader(x)];
    }

    int count() const { return num; }
};

#line 2 "library/datastructure/unionfind/PotentialUnionFind.hpp"
template <typename AbelGroup> class PotentialUnionFind {
    using T = typename AbelGroup::value_type;
    int n, num;
    std::vector<int> sz, parent;
    std::vector<T> potential; // parent[x] を基準とした時の x の値
  public:
    PotentialUnionFind() = default;
    PotentialUnionFind(int n)
        : n(n), num(n), sz(n, 1), parent(n, 0),
          potential(n, AbelGroup::unit()) {
        assert(AbelGroup::commute);
        std::iota(parent.begin(), parent.end(), 0);
    }

    std::pair<int, T> from_root(int x) {
        if (x == parent[x])
            return {x, AbelGroup::unit()};
        auto [r, add] = from_root(parent[x]);
        parent[x] = r;
        AbelGroup::Rchop(potential[x], add);
        return {r, potential[x]};
    }

    int leader(int x) { return from_root(x).first; }

    bool same(int x, int y) {
        assert(0 <= x and x < n and 0 <= y and y < n);
        return leader(x) == leader(y);
    }

    bool merge(int x, int y, T d) {
        // potential[y]-potential[x]=d にする
        // 矛盾する場合は変更はせず false を返す
        assert(0 <= x and x < n and 0 <= y and y < n);
        auto [rx, dx] = from_root(x);
        auto [ry, dy] = from_root(y);
        AbelGroup::Rchop(d, dx);
        AbelGroup::Rchop(d, AbelGroup::inverse(dy));
        if (rx == ry)
            return d == AbelGroup::unit();
        if (sz[rx] < sz[ry]) {
            std::swap(rx, ry);
            d = AbelGroup::inverse(d);
        }
        sz[rx] += sz[ry];
        parent[ry] = rx;
        potential[ry] = d;
        num--;
        return true;
    }

    std::optional<T> diff(int x, int y) {
        // x を基準とする
        auto [rx, dx] = from_root(x);
        auto [ry, dy] = from_root(y);
        if (rx != ry)
            return std::nullopt;
        return AbelGroup::op(dy, AbelGroup::inverse(dx));
    }

    int size(const int x) {
        assert(0 <= x and x < n);
        return sz[leader(x)];
    }

    int count() const { return num; }
};