Weighted DSU
(datastructure/weighted_dsu.hpp)

• View this file on GitHub
• Last update: 2026-10-26 01:30:50+09:00
• Include: #include "datastructure/weighted_dsu.hpp"

概要

• 各頂点にポテンシャルを持つ Union-Find。頂点間のポテンシャルの差を制約として追加し、差を求める。

使い方

• weighted_dsu<G>(n): $n$ 頂点 $0$ 辺の無向グラフを作る。G は S, op, e, inv を static メンバとして持つ群。
• merge(a, b, w): diff(a, b) == w となるように辺 $(a, b)$ を追加し、既存の制約と矛盾しないかを返す。矛盾するときは何もしない。
• diff(a, b): 頂点 a から頂点 b へのポテンシャルの差 op(inv(p[a]), p[b]) を返す。
• same(a, b): 頂点 a, b が連結かを返す。
• leader(a): 頂点 a の属する連結成分の代表元を返す。
• size(a): 頂点 a の属する連結成分のサイズを返す。
• groups(): 連結成分ごとの頂点のリストを返す。

制約

• $0 \le a, b < n$
• diff: same(a, b)
• S は == で比較できる

計算量

• merge, diff, same, leader, size: 償却 $\mathcal{O}(\alpha(n))$
• groups: $\mathcal{O}(n)$

Depends on

• Algebra: Concept (algebra/concept.hpp)

Verified with

• test/library-checker/unionfind_with_potential.test.cpp
• test/library-checker/unionfind_with_potential_non_commutative_group.test.cpp
• test/unit/weighted_dsu.test.cpp

Code

#pragma once
#include <algorithm>
#include <cassert>
#include <utility>
#include <vector>

#include "algebra/concept.hpp"

// potential p[v] of each vertex; diff(a, b) = op(inv(p[a]), p[b])
template <group G> struct weighted_dsu {
public:
    using S = typename G::S;
    weighted_dsu() : weighted_dsu(0) {
    }
    explicit weighted_dsu(int n) : _n(n), parent_or_size(n, -1), pot(n, G::e()) {
    }
    // make diff(a, b) == w; returns false if it contradicts the existing constraints
    bool merge(int a, int b, S w) {
        assert(0 <= a && a < _n);
        assert(0 <= b && b < _n);
        int x = leader(a), y = leader(b);
        if (x == y) return diff(a, b) == w;
        // p[y] relative to x: op(inv(p[a]), op(p[y], p[b])) == w
        S py = G::op(G::op(pot[a], w), G::inv(pot[b]));
        if (-parent_or_size[x] < -parent_or_size[y]) {
            std::swap(x, y);
            py = G::inv(py);
        }
        parent_or_size[x] += parent_or_size[y];
        parent_or_size[y] = x;
        pot[y] = py;
        return true;
    }
    S diff(int a, int b) {
        assert(same(a, b));
        return G::op(G::inv(pot[a]), pot[b]);
    }
    bool same(int a, int b) {
        assert(0 <= a && a < _n);
        assert(0 <= b && b < _n);
        return leader(a) == leader(b);
    }
    int leader(int a) {
        assert(0 <= a && a < _n);
        if (parent_or_size[a] < 0) return a;
        int p = parent_or_size[a];
        int r = leader(p);
        if (p != r) pot[a] = G::op(pot[p], pot[a]);
        return parent_or_size[a] = r;
    }
    int size(int a) {
        assert(0 <= a && a < _n);
        return -parent_or_size[leader(a)];
    }
    std::vector<std::vector<int>> groups() {
        std::vector<int> leader_buf(_n), group_size(_n);
        for (int i = 0; i < _n; i++) {
            leader_buf[i] = leader(i);
            group_size[leader_buf[i]]++;
        }
        std::vector<std::vector<int>> result(_n);
        for (int i = 0; i < _n; i++) result[i].reserve(group_size[i]);
        for (int i = 0; i < _n; i++) result[leader_buf[i]].push_back(i);
        result.erase(std::remove_if(result.begin(), result.end(),
                                    [&](const std::vector<int>& v) { return v.empty(); }),
                     result.end());
        return result;
    }

private:
    int _n;
    std::vector<int> parent_or_size;
    std::vector<S> pot;  // relative to parent
};

#line 2 "datastructure/weighted_dsu.hpp"
#include <algorithm>
#include <cassert>
#include <utility>
#include <vector>

#line 2 "algebra/concept.hpp"
#include <concepts>
template <class M>
concept semigroup = requires(typename M::S a, typename M::S b) {
    { M::op(a, b) } -> std::same_as<typename M::S>;
};
template <class M>
concept monoid = semigroup<M> && requires {
    { M::e() } -> std::same_as<typename M::S>;
};
template <class G>
concept group = monoid<G> && requires(typename G::S a) {
    { G::inv(a) } -> std::same_as<typename G::S>;
};
template <class A>
concept action = requires(typename A::F f, typename A::F g) {
    { A::composition(f, g) } -> std::same_as<typename A::F>;
    { A::id() } -> std::same_as<typename A::F>;
    { A::composition_commutative } -> std::convertible_to<bool>;
};
template <class A>
concept acted_monoid = monoid<A> && action<A> && requires(typename A::S x, typename A::F f) {
    { A::mapping(f, x) } -> std::same_as<typename A::S>;
};
template <class R>
concept semiring = requires(typename R::S a, typename R::S b) {
    { R::add(a, b) } -> std::same_as<typename R::S>;
    { R::zero() } -> std::same_as<typename R::S>;
    { R::mul(a, b) } -> std::same_as<typename R::S>;
    { R::one() } -> std::same_as<typename R::S>;
};
#line 8 "datastructure/weighted_dsu.hpp"

// potential p[v] of each vertex; diff(a, b) = op(inv(p[a]), p[b])
template <group G> struct weighted_dsu {
public:
    using S = typename G::S;
    weighted_dsu() : weighted_dsu(0) {
    }
    explicit weighted_dsu(int n) : _n(n), parent_or_size(n, -1), pot(n, G::e()) {
    }
    // make diff(a, b) == w; returns false if it contradicts the existing constraints
    bool merge(int a, int b, S w) {
        assert(0 <= a && a < _n);
        assert(0 <= b && b < _n);
        int x = leader(a), y = leader(b);
        if (x == y) return diff(a, b) == w;
        // p[y] relative to x: op(inv(p[a]), op(p[y], p[b])) == w
        S py = G::op(G::op(pot[a], w), G::inv(pot[b]));
        if (-parent_or_size[x] < -parent_or_size[y]) {
            std::swap(x, y);
            py = G::inv(py);
        }
        parent_or_size[x] += parent_or_size[y];
        parent_or_size[y] = x;
        pot[y] = py;
        return true;
    }
    S diff(int a, int b) {
        assert(same(a, b));
        return G::op(G::inv(pot[a]), pot[b]);
    }
    bool same(int a, int b) {
        assert(0 <= a && a < _n);
        assert(0 <= b && b < _n);
        return leader(a) == leader(b);
    }
    int leader(int a) {
        assert(0 <= a && a < _n);
        if (parent_or_size[a] < 0) return a;
        int p = parent_or_size[a];
        int r = leader(p);
        if (p != r) pot[a] = G::op(pot[p], pot[a]);
        return parent_or_size[a] = r;
    }
    int size(int a) {
        assert(0 <= a && a < _n);
        return -parent_or_size[leader(a)];
    }
    std::vector<std::vector<int>> groups() {
        std::vector<int> leader_buf(_n), group_size(_n);
        for (int i = 0; i < _n; i++) {
            leader_buf[i] = leader(i);
            group_size[leader_buf[i]]++;
        }
        std::vector<std::vector<int>> result(_n);
        for (int i = 0; i < _n; i++) result[i].reserve(group_size[i]);
        for (int i = 0; i < _n; i++) result[leader_buf[i]].push_back(i);
        result.erase(std::remove_if(result.begin(), result.end(),
                                    [&](const std::vector<int>& v) { return v.empty(); }),
                     result.end());
        return result;
    }

private:
    int _n;
    std::vector<int> parent_or_size;
    std::vector<S> pot;  // relative to parent
};

Back to top page