CCF 201412-4 最优灌溉

#include "iostream"
#include "limits.h"
#include "algorithm"
#define N 1001

using namespace std;

int graph[N][N]; // asumme index start from 1
int V, E; // vertex number, edge number

int minIdx(int key[], bool mstSet[])
{
  int min = INT_MAX, min_idx = -1;
  for(int i=1; i<=V; i++)
  {
    if(!mstSet[i] && key[i] < min)
    {
      min = key[i];
      min_idx = i;
    }
  }
  return min_idx;
}

int primMST()
{
  int key[N];
  for(int i=1; i<=V; i++) key[i] = INT_MAX;
  
  bool mstSet[N]; // indicate wether in minimum spanning tree set
  for(int i=1; i<=V; i++) mstSet[i] = false;
  key[1] = 0;
  
  int cost = 0;
  for(int i=1; i<=V; i++)
  {
    int u = minIdx(key, mstSet);
    mstSet[u] = true;
    cost += key[u];
    for(int v=1; v<=V; v++)
    {
      if(graph[u][v]<INT_MAX && !mstSet[v] && graph[u][v]<key[v])
        key[v] = graph[u][v];
    }
  }
  return cost;
}

int init()
{
  for(int i=1; i<=V; i++)
    for(int j=1; j<=V; j++)
      graph[i][j] = INT_MAX;
}

int main() 
{
  cin >> V >> E;
  init();
  for(int i=0; i<E; i++)
  {
    int e1, e2, w;
    cin >> e1 >> e2 >> w;
    graph[e1][e2] = graph[e2][e1] = min(w, graph[e1][e2]);
  }
  cout << primMST();
}