UF_ComponentSum
(DataStructure/UF_ComponentSum.hpp)

View this file on GitHub
Last update: 2025-05-27 22:26:13+09:00
Include: #include "DataStructure/UF_ComponentSum.hpp"

概要

連結成分に情報（sum）を持たせた上で、merge を行う。merge 後の連結成分の sum は、merge 前の 2 連結成分の sum を基に計算を行う。

テンプレート引数

template <typename T, merge_function<T> f>

連結成分に持たせる sum の型と、merge のときの sum の計算を行う関数を渡す。T は可換でないといけない。また、ラムダ関数をテンプレート引数として使用することはできないため、外部引数を渡すこと。

コンストラクタ

    UF_ComponentSum() = default;
    UF_ComponentSum(const std::vector<T>& init) : 
        UnionFind((int)init.size()), sum_(init) {}

各連結成分の sum の情報は、vector として保持しておくが、その初期値を引数として渡す。

`merge_function`

template <typename T>
using merge_function = void(*)(T& component_sum1, T component_sum2);

2 つの引数両方に、merge の計算のときのオペランド（つまり、被 merge 連結成分の sum）を渡す。第一引数の方の sum を、merge 後の連結成分の sum として再利用することに注意。

機能

merge

bool merge (const int& u, const int& v)

u の属する連結成分と v の属する連結成分を merge する。このとき、merge 後に生成される連結成分の sum も同時に計算する。

sum

T sum (const int& v) const 

v が属する連結成分の sum を返す。

Depends on

DataStructure/UnionFind.hpp

Code

#ifndef UF_ComponentSum_HPP
#define UF_ComponentSum_HPP

#include <vector>
#include <cassert>

#include "UnionFind.hpp"

template <typename Op>
class UF_ComponentSum : public UnionFind { // 継承
    int N; // 頂点数
    Op op;

public:
                                                             // UF を初期化
    UF_ComponentSum(int N, const Op& op) : N(N), op(op), UnionFind(N) {}

    bool merge(int u, int v) { // 上書き
        int ru = root(u), rv = root(v);
        int r = UnionFind::merge(u, v);
        bool merged = (r != ru) || (r != rv);
        if (merged) {
            if (r == ru) {
                op(r, rv);
            } else { // r == pv
                op(r, ru);
            }
        }
        return merged;
    }
};

#endif // UF_ComponentSum_HPP

#line 1 "DataStructure/UF_ComponentSum.hpp"



#include <vector>
#include <cassert>

#line 1 "DataStructure/UnionFind.hpp"



/*
verify
・https://judge.yosupo.jp/problem/unionfind
*/

#line 10 "DataStructure/UnionFind.hpp"
#include <algorithm>
#include <numeric>

class UnionFind {
private:
    int N_;
    std::vector<int> parent_;
    std::vector<int> size_;

public:
    UnionFind() = default;
    UnionFind(int N) : N_(N), parent_(N), size_(N, 1) {
        std::iota(parent_.begin(), parent_.end(), 0);
    }

    int root(int v) {
        static std::vector<int> tmp;
        while (parent_[v] != v) { // 根まで辿っていく
            tmp.push_back(v);
            v = parent_[v];
        }
        while (!tmp.empty()) { // 経路圧縮
            parent_[tmp.back()] = v;
            tmp.pop_back();
        }
        return v;
    }

    bool same(int u, int v) {
        return root(u) == root(v);
    }

    int merge(int u, int v) {
        int root_u = root(u), root_v = root(v);

        if (root_u == root_v) {
            return root_u; // 根が同じなら、既に同じ集合
        }

        if (size_[root_u] < size_[root_v]) {
            std::swap(root_u, root_v);
        }
        size_[root_u] += size_[root_v];
        parent_[root_v] = root_u;
        return root_u;
    }

    int size(int v) {
        return size_[root(v)];
    }

    std::vector<std::vector<int>> groups() {
        std::vector<std::vector<int>> res(N_);
        for (int i = 0; i < N_; i++) {
            res[root(i)].push_back(i);
        }
        res.erase(std::remove_if(res.begin(), res.end(), [](const std::vector<int>& vec){ return vec.empty(); }), res.end());
        return res;
    }
};


#line 8 "DataStructure/UF_ComponentSum.hpp"

template <typename Op>
class UF_ComponentSum : public UnionFind { // 継承
    int N; // 頂点数
    Op op;

public:
                                                             // UF を初期化
    UF_ComponentSum(int N, const Op& op) : N(N), op(op), UnionFind(N) {}

    bool merge(int u, int v) { // 上書き
        int ru = root(u), rv = root(v);
        int r = UnionFind::merge(u, v);
        bool merged = (r != ru) || (r != rv);
        if (merged) {
            if (r == ru) {
                op(r, rv);
            } else { // r == pv
                op(r, ru);
            }
        }
        return merged;
    }
};

UF_ComponentSum (DataStructure/UF_ComponentSum.hpp)

概要