Weighted Union Find
(data-structure/unionfind/weighted_union_find.hpp)

• View this file on GitHub
• Last update: 2024-01-07 20:49:49+09:00
• Include: #include "data-structure/unionfind/weighted_union_find.hpp"

Description

重み付き union find は，union find の操作に加えて，同じ集合に属する他の要素に対する要素の相対的な重みを管理する．

空間計算量: $O(n)$

Operations

• WeightedUnionFind(int n)
  • サイズnの重み付き union find を構築する．
  • 時間計算量: $O(n)$
• int find(int x)
  • $x$ が属する木の根を返す
  • 時間計算量: $\mathrm{amortized}\ O(\alpha(n))$
• T weight(int x)
  • 木の根に対する $x$ の重みを返す
  • 時間計算量: $\mathrm{amortized}\ O(\alpha(n))$
• void unite(int x, int y, T w)
  • $x$ が属する集合と $y$ が属する集合を $weight(y) - weight(x) = w$ となるように連結する
  • 時間計算量: $\mathrm{amortized}\ O(\alpha(n))$
• bool same(int x, int y)
  • $x$ と $y$ が同じ集合に属するかを判定する
  • 時間計算量: $\mathrm{amortized}\ O(\alpha(n))$
• T diff(int x, int y)
  • $x$ に対する $y$ の重み，すなわち $weight(y) - weight(x)$ を返す
  • 時間計算量: $\mathrm{amortized}\ O(\alpha(n))$
• int size(int x)
  • $x$ が属する集合の大きさを返す
  • 時間計算量: $\mathrm{amortized}\ O(\alpha(n))$

Verified with

• test/aoj/DSL_1_B.test.cpp

Code

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

template <typename T>
class WeightedUnionFind {
   public:
    WeightedUnionFind() = default;
    explicit WeightedUnionFind(int n) : data(n, -1), ws(n) {}

    int find(int x) {
        if (data[x] < 0) return x;
        int r = find(data[x]);
        ws[x] += ws[data[x]];
        return data[x] = r;
    }

    T weight(int x) {
        find(x);
        return ws[x];
    }

    bool unite(int x, int y, T w) {
        w += weight(x);
        w -= weight(y);
        x = find(x);
        y = find(y);
        if (x == y) return false;
        if (data[x] > data[y]) {
            std::swap(x, y);
            w = -w;
        }
        data[x] += data[y];
        data[y] = x;
        ws[y] = w;
        return true;
    }

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

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

    int size(int x) { return -data[find(x)]; }

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

#line 2 "data-structure/unionfind/weighted_union_find.hpp"
#include <algorithm>
#include <vector>

template <typename T>
class WeightedUnionFind {
   public:
    WeightedUnionFind() = default;
    explicit WeightedUnionFind(int n) : data(n, -1), ws(n) {}

    int find(int x) {
        if (data[x] < 0) return x;
        int r = find(data[x]);
        ws[x] += ws[data[x]];
        return data[x] = r;
    }

    T weight(int x) {
        find(x);
        return ws[x];
    }

    bool unite(int x, int y, T w) {
        w += weight(x);
        w -= weight(y);
        x = find(x);
        y = find(y);
        if (x == y) return false;
        if (data[x] > data[y]) {
            std::swap(x, y);
            w = -w;
        }
        data[x] += data[y];
        data[y] = x;
        ws[y] = w;
        return true;
    }

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

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

    int size(int x) { return -data[find(x)]; }

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

Back to top page