Weighted Union-Find
(lib/data_structure/weighted_union_find.hpp)

• View this file on GitHub
• Last update: 2024-11-18 03:54:10+09:00
• Include: #include "lib/data_structure/weighted_union_find.hpp"

Weighted Union Find

使い方

• WeightedUnionFind(N): 要素数 N の Weighted Union Find を構築します。
• root(x): 集合 x の根を取得します。
• diff(x, y): x と y の重みの差を取得します。
• unite(x, y, w): 集合 x と y について、重みが w として併合します。
• same(x, y): x と y の属する集合が等しいかどうかを判定します。
• size(x): x の属する集合の大きさを取得します。

計算量

• WeightedUnionFind(N): $\mathrm{O}(N)$
• root(x): 償却 $\mathrm{O}(\alpha(N))$ (アッカーマンの逆関数)
• diff(x, y): 償却 $\mathrm{O}(\alpha(N))$
• unite(x, y, w): 償却 $\mathrm{O}(\alpha(N))$
• same(x, y): 償却 $\mathrm{O}(\alpha(N))$
• size(x): 償却 $\mathrm{O}(\alpha(N))$

Verified with

• test/aoj/dsl/dsl_1_b.test.cpp
• test/library_checker/data_structure/unionfind_with_potential.test.cpp
• test/library_checker/data_structure/unionfind_with_potential_non_commutative_group.test.cpp

Code

#pragma once

/**
 * @brief Weighted Union-Find
 * @docs docs/data_structure/weighted_union_find.md
 */

#include <vector>
#include <cassert>

template <typename T>
struct WeightedUnionFind{
    int V;
    std::vector<int> par;
    std::vector<T> diff_weight;
    T ex;

    WeightedUnionFind(const int N, const T &e = 0) : V(N), par(N), diff_weight(N), ex(e){
        for(int i = 0; i < N; ++i){
            par[i] = -1;
            diff_weight[i] = e;
        }
    }

    int root(const int x){
        assert(0 <= x && x < V);
        if(par[x] < 0){
            return x;
        }
        int px = par[x];
        int rx = root(par[x]);
        diff_weight[x] = diff_weight[x] + diff_weight[px];
        return par[x] = rx;
    }

    T weight(const int x){
        root(x);
        return diff_weight[x];
    }

    T diff(const int x, const int y){
        return weight(x) - weight(y);
    }

    void unite(const int x, const int y, const T &w){
        int tx = x, ty = y;
        T tw = w;
        int rx = root(x), ry = root(y);
        if(rx == ry) return;

        if(par[rx] < par[ry]){
            std::swap(rx, ry);
            std::swap(tx, ty);
            tw = ex - tw;
        }

        par[ry] = par[rx] + par[ry];
        par[rx] = ry;
        diff_weight[rx] = ex - diff_weight[tx] + tw + diff_weight[ty];
    }

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

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

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

/**
 * @brief Weighted Union-Find
 * @docs docs/data_structure/weighted_union_find.md
 */

#include <vector>
#include <cassert>

template <typename T>
struct WeightedUnionFind{
    int V;
    std::vector<int> par;
    std::vector<T> diff_weight;
    T ex;

    WeightedUnionFind(const int N, const T &e = 0) : V(N), par(N), diff_weight(N), ex(e){
        for(int i = 0; i < N; ++i){
            par[i] = -1;
            diff_weight[i] = e;
        }
    }

    int root(const int x){
        assert(0 <= x && x < V);
        if(par[x] < 0){
            return x;
        }
        int px = par[x];
        int rx = root(par[x]);
        diff_weight[x] = diff_weight[x] + diff_weight[px];
        return par[x] = rx;
    }

    T weight(const int x){
        root(x);
        return diff_weight[x];
    }

    T diff(const int x, const int y){
        return weight(x) - weight(y);
    }

    void unite(const int x, const int y, const T &w){
        int tx = x, ty = y;
        T tw = w;
        int rx = root(x), ry = root(y);
        if(rx == ry) return;

        if(par[rx] < par[ry]){
            std::swap(rx, ry);
            std::swap(tx, ty);
            tw = ex - tw;
        }

        par[ry] = par[rx] + par[ry];
        par[rx] = ry;
        diff_weight[rx] = ex - diff_weight[tx] + tw + diff_weight[ty];
    }

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

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

Back to top page