graph/kruskal.hpp

View this file on GitHub
Last update: 2021-05-27 03:39:22+09:00
Include: #include "graph/kruskal.hpp"

Depends on

datastructure/UnionFind.hpp

Verified with

test/aoj-grl-2-a.test.cpp

Code

#include<queue>
#include<vector>
#include<algorithm>
#include "../datastructure/UnionFind.hpp"
template<typename T>
T kruskal(std::vector<std::vector<std::pair<int, T>>> &g) {
  int n = g.size();
  std::vector<std::pair<T, std::pair<int, int>>> g2;
  for(int v=0; v<n; v++) {
    for(auto[u, w] : g[v]) {
      g2.push_back({w, {v, u}});
    }
  }
  sort(g2.begin(), g2.end());
  UnionFind uf(n);
  T ans = 0;
  for(int i=0; i<(int)g2.size(); i++){
    T w = g2[i].first;
    auto[v, u] = g2[i].second;
    if(!uf.issame(v, u)) {
      uf.merge(v, u);
      ans += w;
    }
  }
  return ans;
}

#line 1 "graph/kruskal.hpp"
#include<queue>
#include<vector>
#include<algorithm>
#line 2 "datastructure/UnionFind.hpp"
#include <cassert>
#include <iostream>
struct UnionFind{
  std::vector<int> par;
  std::vector<int> siz;

  UnionFind(int sz_): par(sz_), siz(sz_) {
    for(int i=0; i<sz_; ++i) par[i] = i, siz[i] = 1;
  }

  void init(int sz_){
    par.resize(sz_);
    siz.resize(sz_);
    for(int i=0; i<sz_; ++i) par[i] = i, siz[i] = 1;
  }

  int root(int x){
    if(x == par[x]) return x;
    return par[x] = root(par[x]);
  }

  bool merge(int x, int y){
    x = root(x), y = root(y);
    if(x == y) return false;
    if(siz[x] < siz[y]) std::swap(x, y);
    siz[x] += siz[y];
    par[y] = x;
    return true;
  }

  bool issame(int x, int y){
    return root(x) == root(y);
  }

  int size(int x){
    return siz[root(x)];
  }
};
#line 5 "graph/kruskal.hpp"
template<typename T>
T kruskal(std::vector<std::vector<std::pair<int, T>>> &g) {
  int n = g.size();
  std::vector<std::pair<T, std::pair<int, int>>> g2;
  for(int v=0; v<n; v++) {
    for(auto[u, w] : g[v]) {
      g2.push_back({w, {v, u}});
    }
  }
  sort(g2.begin(), g2.end());
  UnionFind uf(n);
  T ans = 0;
  for(int i=0; i<(int)g2.size(); i++){
    T w = g2[i].first;
    auto[v, u] = g2[i].second;
    if(!uf.issame(v, u)) {
      uf.merge(v, u);
      ans += w;
    }
  }
  return ans;
}