重み付きUnion-Find木

このデータ構造の実装には、 けんちょんさんの解説 を大変参考にさせていただきました。

    
/*
        この重み付きunion_find木で出来ること
        例)weight_union_find tree(N);などで宣言(N:頂点数)
        tree.same(x,y):要素xと要素yが同じグループに属するか判定
        tree.unite(x,y,w):weight(y)=weight(x)+wとなるように木を併合
        tree.diff(x,y):weight(y)-weight(x)を返す
        tree.size(x):要素xが属する木の頂点数を返す
        */
struct weight_union_find
{
    vector<long long> par;      //親の番号 
    vector<long long> rank;     //木の深さ(根のランクは0)
    vector<long long> siz;      //要素xが根なら木の頂点数を格納する
    vector<long long> par_diff; //親と子のコストの差
    //初期化子リストを用いた初期化
    weight_union_find(long long N) : par(N), rank(N), siz(N), par_diff(N)
    {
        for (long long i = 0; i < N; i++)
        {
            par[i] = i;
            rank[i] = 0;
            siz[i] = 1;
            par_diff[i] = 0;
        }
    }
    //要素xが所属する木の根を再帰的に発見する
    long long root(long long x)
    {
        if (par[x] == x)
            return x;
        else
        {
            //経路圧縮+累積和
            long long ret = root(par[x]);
            par_diff[x] += par_diff[par[x]];
            return par[x] = ret;
        }
    }
    //経路圧縮+コストを返す
    long long weight(long long x)
    {
        root(x);
        return par_diff[x];
    }
    //yのコスト-xのコスト
    long long diff(long long x, long long y)
    {
        return weight(y) - weight(x);
    }
    //weight(y)=weight(x)+wとなるように木を併合
    void unite(long long x, long long y, long long w)
    {
        long long rx = root(x);
        long long ry = root(y);
        w += weight(x);
        w -= weight(y);
        if (rx == ry)
            return; //同じ木に属してたらそのまま
        if (rank[rx] < rank[ry])
        {
            swap(rx, ry);
            siz[ry] = siz[rx] + siz[ry];
            w = -w; //符号反転
        }
        else
        {
            siz[rx] = siz[rx] + siz[ry];
        }
        if (rank[rx] == rank[ry])
            rank[rx]++;
        par[ry] = rx;
        par_diff[ry] = w;
    }
    //要素xが属する木と要素yが属する木が同じならtrueを返す
    bool same(long long x, long long y)
    {
        return root(x) == root(y);
    }
    //要素xが属する木の頂点数を返す
    long long size(long long x)
    {
        return siz[root(x)];
    }
};
    
    
© 2020 kacho65535