Union Find With Potential (data_structure/unionfind_with_potential.hpp)

• View this file on GitHub
• View document part on GitHub
• Last update: 2025-04-27 17:12:50+09:00
• Include: #include "data_structure/unionfind_with_potential.hpp"

集合内の要素間の重みの差を持ちながら、互いに素な集合を管理します。

コンストラクタ

UnionfindWithPotential(int n)

説明

n 個の要素を持つ Weighted Union-Find を構築します。最初は各要素がそれぞれ異なる集合に属しています。

• n: 要素数。

計算量

$O(N)$

merge

int merge(int a, int b, T W)

説明

要素 a と要素 b を同じ集合にマージし、$b - a = W$ という関係を確立します。 もし a と b が既に同じ集合に属している場合、この操作は $b - a = W$ という関係が既存の関係と矛盾しないかチェックする役割も果たしますが、この実装では矛盾の検出は行わず、単にマージ操作を行います。

• a: 要素のインデックス (0-indexed)。
• b: 要素のインデックス (0-indexed)。
• W: b - a の重み差分。
• 戻り値: マージ後の集合の代表元のインデックス。

制約

$0 \le a < n$
$0 \le b < n$

計算量

$O(\alpha(N))$

same

bool same(int a, int b)

説明

要素 a と要素 b が同じ集合に属しているかを判定します。

• a: 要素のインデックス (0-indexed)。
• b: 要素のインデックス (0-indexed)。
• 戻り値: a と b が同じ集合に属していれば true、そうでなければ false。

制約

$0 \le a < n$
$0 \le b < n$

計算量

$O(\alpha(N))$

weight

T weight(int a, int b)

説明

要素 a と要素 b が同じ集合に属している場合、$b - a$ の重み差分を返します。 異なる集合に属している場合の挙動は未定義です。

• a: 要素のインデックス (0-indexed)。
• b: 要素のインデックス (0-indexed)。
• 戻り値: b - a の重み差分。

制約

$0 \le a < n$
$0 \le b < n$

計算量

$O(\alpha(N))$

leader

std::pair<int, T> leader(int a)

説明

要素 a が属する集合の代表元と、a からその代表元までの重み差分を返します。 具体的には、代表元を $r$ とすると、戻り値は {r, r - a} となります。

• a: 要素のインデックス (0-indexed)。
• 戻り値: {代表元のインデックス, 代表元 - a の重み差分} のペア。

制約

$0 \le a < n$

計算量

$O(\alpha(N))$

size

int size(int a)

説明

要素 a が属する集合の要素数を返します。

• a: 要素のインデックス (0-indexed)。
• 戻り値: 要素 a が属する集合のサイズ。

制約

$0 \le a < n$

計算量

$O(\alpha(N))$

Verified with

• test/DSL_1_B.test.cpp
• test/unionfind_with_potential.test.cpp

Code

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

template <typename T = int64_t>
struct UnionFindWithPotential {
    int n;
    std::vector<int> v;
    std::vector<T> w;

    UnionFindWithPotential() : n(0) {}
    explicit UnionFindWithPotential(int n) : n(n), v(n, -1), w(n, 0) {}

    int merge(int a, int b, T W) {
        auto [x, sx] = leader(a);
        auto [y, sy] = leader(b);
        if (x == y) return x;
        W = -W;
        if (-v[x] < -v[y]) {
            std::swap(x, y);
            std::swap(sx, sy);
            W = -W;
        }
        v[x] += v[y];
        v[y] = x;
        w[y] = (sx - sy) - W;
        return x;
    }

    bool same(int a, int b) {
        return leader(a).first == leader(b).first;
    }

    // return b - a
    T weight(int a, int b) {
        return leader(b).second - leader(a).second;
    }

    std::pair<int, T> leader(int a, T sum = 0) {
        if (v[a] < 0) return {a, sum};
        return leader(v[a], sum + w[a]);
    }

    int size(int a) {
        return -v[leader(a).first];
    }
};

#line 2 "data_structure/unionfind_with_potential.hpp"
#include <algorithm>
#include <vector>

template <typename T = int64_t>
struct UnionFindWithPotential {
    int n;
    std::vector<int> v;
    std::vector<T> w;

    UnionFindWithPotential() : n(0) {}
    explicit UnionFindWithPotential(int n) : n(n), v(n, -1), w(n, 0) {}

    int merge(int a, int b, T W) {
        auto [x, sx] = leader(a);
        auto [y, sy] = leader(b);
        if (x == y) return x;
        W = -W;
        if (-v[x] < -v[y]) {
            std::swap(x, y);
            std::swap(sx, sy);
            W = -W;
        }
        v[x] += v[y];
        v[y] = x;
        w[y] = (sx - sy) - W;
        return x;
    }

    bool same(int a, int b) {
        return leader(a).first == leader(b).first;
    }

    // return b - a
    T weight(int a, int b) {
        return leader(b).second - leader(a).second;
    }

    std::pair<int, T> leader(int a, T sum = 0) {
        if (v[a] < 0) return {a, sum};
        return leader(v[a], sum + w[a]);
    }

    int size(int a) {
        return -v[leader(a).first];
    }
};

Back to top page