重み付きUnionFind
(datastructure/unionfind/weightedunionfind.hpp)

• View this file on GitHub
• Last update: 2020-08-25 17:20:34+09:00
• Include: #include "datastructure/unionfind/weightedunionfind.hpp"
• Link: https://qiita.com/drken/items/cce6fc5c579051e64fab

概要

連結成分とノード間の距離を管理するデータ構造.

計算量

$O(\alpha (n))$

使用例

• WeightedUnionFind uf(V) : $V$頂点の重み付きUnionFindを定義.
• uf.unite(x, y, w) : $x$と$y$を連結にする.
  • このとき, $\mathrm{weight}(y) = \mathrm{weight}(x) + w$となるようにする.
• uf.same(x, y) : $x$と$y$が連結かを返す.
• uf.size(x) : $x$の属する連結成分の大きさを返す.
• uf.size() : 連結成分の個数を返す.

Verified with

• test/aoj/1330.test.cpp
• test/aoj/DSL_1_B.test.cpp
• test/yosupo/unionfind.weightedunionfind.test.cpp

Code

/**
* @brief 重み付きUnionFind
* @docs docs/datastructure/unionfind/weightedunionfind.md
* @see https://qiita.com/drken/items/cce6fc5c579051e64fab
*/

template <typename T>
struct WeightedUnionFind {
    int sz;
    vector<int> parent;
    vector<T> diff_weight;
    WeightedUnionFind(int sz)
        : sz(sz), parent(sz, -1), diff_weight(sz, T(0)) {}
    // weight(y) = weight(x) + w
    bool unite(int x, int y, T w) {
        w += weight(x) - weight(y);
        if ((x = find(x)) != (y = find(y))) {
            if (parent[y] < parent[x]) swap(x, y), w = -w;
            --sz;
            parent[x] += parent[y];
            parent[y]      = x;
            diff_weight[y] = w;
        }
        return true;
    }
    int find(int x) {
        if (parent[x] < 0) return x;
        int ret = find(parent[x]);
        diff_weight[x] += diff_weight[parent[x]];
        return parent[x] = ret;
    }
    T weight(int x) {
        find(x);
        return diff_weight[x];
    }
    T diff(int x, int y) { return weight(y) - weight(x); }
    bool same(int x, int y) { return find(x) == find(y); }
    int size(int x) { return -parent[find(x)]; }
    int size() { return sz; }
};

#line 1 "datastructure/unionfind/weightedunionfind.hpp"
/**
* @brief 重み付きUnionFind
* @docs docs/datastructure/unionfind/weightedunionfind.md
* @see https://qiita.com/drken/items/cce6fc5c579051e64fab
*/

template <typename T>
struct WeightedUnionFind {
    int sz;
    vector<int> parent;
    vector<T> diff_weight;
    WeightedUnionFind(int sz)
        : sz(sz), parent(sz, -1), diff_weight(sz, T(0)) {}
    // weight(y) = weight(x) + w
    bool unite(int x, int y, T w) {
        w += weight(x) - weight(y);
        if ((x = find(x)) != (y = find(y))) {
            if (parent[y] < parent[x]) swap(x, y), w = -w;
            --sz;
            parent[x] += parent[y];
            parent[y]      = x;
            diff_weight[y] = w;
        }
        return true;
    }
    int find(int x) {
        if (parent[x] < 0) return x;
        int ret = find(parent[x]);
        diff_weight[x] += diff_weight[parent[x]];
        return parent[x] = ret;
    }
    T weight(int x) {
        find(x);
        return diff_weight[x];
    }
    T diff(int x, int y) { return weight(y) - weight(x); }
    bool same(int x, int y) { return find(x) == find(y); }
    int size(int x) { return -parent[find(x)]; }
    int size() { return sz; }
};

Back to top page