WeightedUnionFind
(src/data_structure/weighted_union_find.hpp)

• View this file on GitHub
• Last update: 2024-11-09 01:34:39+09:00
• Include: #include "src/data_structure/weighted_union_find.hpp"

WeightedUnionFind

通常のUnionFindに加えて，

• 頂点に重みを持たせる
• 頂点の重みや頂点間の重みの差の取得
• 重みの情報が矛盾しないかの判定

ができます．

コンストラクタ

WeightedUnionFind<T> uf(int n)

• $n$ 頂点 $0$ 辺の無向グラフを作ります．
• 各頂点の重み $\mathrm{weight_i}$ は $0$ で初期化されます．

計算量

• $O(n)$

merge

bool uf.merge(int a, int b, T w)

辺 $(a, b)$ を追加し， $\mathrm{weight_a} = \mathrm{weight_b} + w$ と設定します．

今までの情報と矛盾が生じない場合は true を返し，矛盾が生じる場合は false を返します．
矛盾が生じるとは merge を行う前の状態において，頂点 $a$ と $b$ が連結かつ$\mathrm{weight_a} - \mathrm{weight_b} \neq w$ であることを指します．

制約

• $0 \leq a < n$
• $0 \leq b < n$

計算量

• 償却 $O(\alpha(n))$

same

bool uf.same(int a, int b)

頂点 $a$ と $b$ が連結かどうかを返します．

制約

• $0 \leq a < n$
• $0 \leq b < n$

計算量

• 償却 $O(\alpha(n))$

leader

int uf.leader(int a)

頂点 $a$ の属する連結成分の代表元を返します．

制約

• $0 \leq a < n$

計算量

• 償却 $O(\alpha(n))$

size

int uf.size(int a)

頂点 $a$ の属する連結成分のサイズを返します．

制約

• $0 \leq a < n$

計算量

• 償却 $O(\alpha(n))$

weight

T uf.weight(int a)

$\mathrm{weight_a} - \mathrm{weight_{leader(a)}}$ を返します．

制約

• $0 \leq a < n$

計算量

• 償却 $O(\alpha(n))$

diff

T uf.diff(int a, int b)

$\mathrm{weight_a} - \mathrm{weight_b}$ を返します．

制約

• $0 \leq a < n$
• $0 \leq b < n$

計算量

• 償却 $O(\alpha(n))$

Depends on

• template (src/template/template.hpp)

Verified with

• verify/aizu_online_judge/dsl/weighted_union_find_trees.test.cpp
• verify/library_checker/data_structure/unionfind_with_potential.test.cpp

Code

#pragma once
#include "../template/template.hpp"
template <typename T>
struct WeightedUnionFind {
    WeightedUnionFind(const int N)
        : n(N), data(N, -1), ws(N, T()) {}
    bool merge(const int a, const int b, T w) {
        assert(0 <= a and a < n);
        assert(0 <= b and b < n);
        w += weight(b) - weight(a);
        int x = leader(a), y = leader(b);
        if(x == y) return w == T();
        if(-data[x] > -data[y]) swap(x, y), w = -w;
        data[y] += data[x];
        data[x] = y;
        ws[x] = w;
        return true;
    }
    bool same(const int a, const int b) {
        assert(0 <= a and a < n);
        assert(0 <= b and b < n);
        return leader(a) == leader(b);
    }
    int leader(const int a) {
        assert(0 <= a and a < n);
        if(data[a] < 0) return a;
        const int r = leader(data[a]);
        ws[a] += ws[data[a]];
        return data[a] = r;
    }
    int size(const int a) {
        assert(0 <= a and a < n);
        return -data[leader(a)];
    }
    T weight(const int a) {
        assert(0 <= a and a < n);
        leader(a);
        return ws[a];
    }
    T diff(const int a, const int b) {
        assert(0 <= a and a < n);
        assert(0 <= b and b < n);
        return weight(a) - weight(b);
    }

   private:
    int n;
    vector<int> data;
    vector<T> ws;
};

#line 2 "src/template/template.hpp"
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using P = pair<long long, long long>;
#define rep(i, a, b) for(long long i = (a); i < (b); ++i)
#define rrep(i, a, b) for(long long i = (a); i >= (b); --i)
constexpr long long inf = 4e18;
struct SetupIO {
    SetupIO() {
        ios::sync_with_stdio(0);
        cin.tie(0);
        cout << fixed << setprecision(30);
    }
} setup_io;
#line 3 "src/data_structure/weighted_union_find.hpp"
template <typename T>
struct WeightedUnionFind {
    WeightedUnionFind(const int N)
        : n(N), data(N, -1), ws(N, T()) {}
    bool merge(const int a, const int b, T w) {
        assert(0 <= a and a < n);
        assert(0 <= b and b < n);
        w += weight(b) - weight(a);
        int x = leader(a), y = leader(b);
        if(x == y) return w == T();
        if(-data[x] > -data[y]) swap(x, y), w = -w;
        data[y] += data[x];
        data[x] = y;
        ws[x] = w;
        return true;
    }
    bool same(const int a, const int b) {
        assert(0 <= a and a < n);
        assert(0 <= b and b < n);
        return leader(a) == leader(b);
    }
    int leader(const int a) {
        assert(0 <= a and a < n);
        if(data[a] < 0) return a;
        const int r = leader(data[a]);
        ws[a] += ws[data[a]];
        return data[a] = r;
    }
    int size(const int a) {
        assert(0 <= a and a < n);
        return -data[leader(a)];
    }
    T weight(const int a) {
        assert(0 <= a and a < n);
        leader(a);
        return ws[a];
    }
    T diff(const int a, const int b) {
        assert(0 <= a and a < n);
        assert(0 <= b and b < n);
        return weight(a) - weight(b);
    }

   private:
    int n;
    vector<int> data;
    vector<T> ws;
};

Back to top page