POJ2377 Bad Cowtractors

POJ2377

#include <iostream>
#include <queue>
using namespace std;

int T, N, M, Q, arr[1005];
priority_queue<pair<int, pair<int, int> > > q;

void ini(void){
  int a, b, c;
  scanf("%d %d", &N, &M);
  while(M--){
    scanf("%d %d %d", &a, &b, &c);
    q.push(make_pair(c, make_pair(a, b) ) );
  }
}

//找頭頭
int par(int x){
  if(arr[x] != x) return arr[x] = par(arr[x]);
  else return x;
}

//這兩個是否同屬一個頭頭
bool check(int a, int b){
  a = par(a);
  b = par(b);
  return (a != b);
}

void sol(void){
  int c, i, j, ans = 0, conn = 0;
  pair<int, int> p;
  
  //初始化每個人的頭頭都是自己
  for(int i = 1; i <= N; i++) arr[i] = i;
  
  //conn 存已選邊的數量，用來檢查是否每點都有
  while(!q.empty() && conn < N){
    c = q.top().first;
    p = q.top().second;
    q.pop();
    
    if(check(arr[p.first], arr[p.second])){
      ans += c;
      conn++;
      
      //將這團的頭頭改為 p.first
      //其實也可以只把這團的頭頭改掉就可了 :D
      //arr[ par(p.second) ] = p.first;
      i = p.second;
      while(arr[i] != i){
        j = arr[i];
        arr[i] = p.first;
        i = j;
      }
      arr[i] = p.first;
      
    }
  }
  
  if(conn == N - 1) printf("%d\n", ans);
  else printf("-1\n");
  
}

int main(void){
  int i, j;

  ini();

  sol();

}