重み付きUnionFind木の実装

重みつきUnionFindは、UnionFind木を理解していればあっという間に実装できます。

と、言うのも、UnionFind木で親をたどっていく操作の途中で、重みも一緒に計算すればいいだけだからです。

ホントにそれだけです。添字がめんどくさいぐらい。

#include <bits/stdc++.h>
using namespace std;
class UnionFind {
public:
  //配列のでかさは調整
  vector<vector<long long int>>par;
  //int par[1000009];
  int siz[1000009];

  // N 頂点の Union-Find を作成
  void init(int N) {
    for (int i = 1; i <= N+10; i++) par.push_back({-1,0}); // 最初は親が無い
    for (int i = 1; i <= N; i++) siz[i] = 1; // 最初はグループの頂点数が 1
  }

  // 頂点 x の根を返す関数
  int root(int x) {
    while (true) {
      if (par[x][0] == -1) break; // 1 個先（親）がなければ、ここが根
      x = par[x][0]; // 1 個先（親）に進む
    }
    return x;
  }

  // 要素 u と v を統合する関数(v+w==u)
  void unite(int u, int v,int w) {
    int RootU = root(u);
    int RootV = root(v);
    if (RootU == RootV) return; // u と v が同グループのときは処理を行わない
    if (siz[RootU] < siz[RootV]) {
      par[RootU][1] = w+kyori_ne(v)+kyori_ne(u)*-1;
      par[RootU][0] = RootV;
      siz[RootV] = siz[RootU] + siz[RootV];
    }
    else {
      par[RootV][1] = w*-1+kyori_ne(u)+kyori_ne(v)*-1;
      par[RootV][0] = RootU;
      siz[RootU] = siz[RootU] + siz[RootV];
    }
  }

  // 要素 u と v が同一のグループかどうかを返す関数　
  bool same(int u, int v) {
    if (root(u) == root(v)) return true;
    return false;
  }
  //その頂点の根との差を出力
  long long int kyori_ne(int x){
    long long int count=0;
    while(true){
      if(par[x][0]==-1)break;
      count+=par[x][1];
      x=par[x][0];
    }
    return count;
  }

  //v-uを出力
  long long int kyori(int u,int v){
    //もし同一グループでなかったら-1を返す
    if(same(u,v)==false)return -1;
    return kyori_ne(u)-kyori_ne(v);
  }

    
  //要素ｎが属しているグループのデカさ
  int dekasa(int n){
    return siz[root(n)];
  }
  void checker(){
    for (int i=0;i<par.size();i++)cout << par[i][0] << " " << par[i][1] << endl;
  }
};