最小全域木の実装

備忘録

最小全域木は、重み付きグラフなどにおいて、連結なグラフにするために最小で重みがいくつになるかです（日本語変かも）。

大きくクラスカル法とプリム法がありますが、ぼくはプリム法の方で実装しました（一般的にはクラスカル法のほうがシンプルなので使われている）。

実装

int purimu(int n,int m,vector<vector<vector<int>>>rotti)
{
  //二次元配列の中身は{行く頂点、距離}
  priority_queue<vector<int>,vector<vector<int>>,greater<vector<int>>>prq;
  //優先度つきキューの中身は{距離、スタート頂点、ゴール頂点}
  set<int>irane{1};
  for (int i=0;i<rotti[0].size();i++)prq.push({rotti[0][i][1],1,rotti[0][i][0]});
  vector<int>hako;
  int ans=0;
  while (prq.size()>0){
    hako=prq.top();
    prq.pop();
    if (irane.count(hako[2])>0)continue;
    ans+=hako[0];
    irane.insert(hako[2]);
    for (int i=0;i<rotti[hako[2]-1].size();i++)prq.push({rotti[hako[2]-1][i][1],hako[2],rotti[hako[2]-1][i][0]});
  }
  if (irane.size()<n)return -1;
  return ans;
}

頂点の数、辺の数、連結関係と重みを格納した三次元配列を引数とし、最小全域木の重みを返します。
もし、どうやっても連結なグラフにできなかった場合、-1を返します。

プリム法について軽く説明すると

最初にスタートする頂点を決める。
そして、優先度付きキューなどに、その頂点から使える辺の長さなどを格納しておく

そして、以下の動作を優先度付きキューが空になるまで続けます。

優先度付きキューから、今使える中で一番短い辺を選ぶ
もしその辺を追加することによって閉路になってしまうのなら以下の動作をスキップする
・その辺を使うことにする。そして、その結果連結となる頂点から使える辺をまた優先度付きキューに追加する

例えば、以下のような実装で最小全域木の重みを出力することができます。

#include <bits/stdc++.h>
using namespace std;
//最小全域木の重みを返す
int purimu(int n,int m,vector<vector<vector<int>>>rotti)
{
  //二次元配列の中身は{行く頂点、距離}
  priority_queue<vector<int>,vector<vector<int>>,greater<vector<int>>>prq;
  //優先度つきキューの中身は{距離、スタート頂点、ゴール頂点}
  set<int>irane{1};
  for (int i=0;i<rotti[0].size();i++)prq.push({rotti[0][i][1],1,rotti[0][i][0]});
  vector<int>hako;
  long long ans=0;
  while (prq.size()>0){
    hako=prq.top();
    prq.pop();
    if (irane.count(hako[2])>0)continue;
    ans+=hako[0];
    irane.insert(hako[2]);
    for (int i=0;i<rotti[hako[2]-1].size();i++)prq.push({rotti[hako[2]-1][i][1],hako[2],rotti[hako[2]-1][i][0]});
  }
  if (irane.size()<n)return -1;
  return ans;
}
int main()
{
  int n,m,a,b,c;
  cin >> n >> m;
  vector<vector<vector<int>>>rotti(n);
  for (int i=0;i<m;i++){
    cin >> a >> b >> c;
    rotti[a-1].push_back({b,c});
    rotti[b-1].push_back({a,c});
  }
  cout << purimu(n,m,rotti) << endl;
  return 0;
}