USACO 3.1 AgriNet

Easy problem with prim MST algorithm.

#include <algorithm>
#include <vector>
#include <iostream>
#include <fstream>
#include <utility>
#include <cmath>
#include <string>
#include <unordered_map>
#include <climits>
#include <sstream>

using namespace std;

int N;
int ret = 0;
vector<bool> minset;
vector<int> parent;
vector<int> dist;
vector<vector<int>> grid;

int main() {
  ifstream fin("agrinet.in");
  ofstream fout("agrinet.out");

  fin >> N;
  grid = vector<vector<int>>(N, vector<int>(N, 0));
  for (int i = 0; i < N; ++i) {
    for (int j = 0; j < N; ++j) {
      fin >> grid[i][j];
    }
  }

  minset = vector<bool>(N, false);
  parent = vector<int>(N, -1);
  dist = vector<int>(N, INT_MAX);
  dist[0] = 0;

  // Iterate N -1 times (since we have initialized dist for vertex 0 already).
  for (int i = 0; i < N - 1; ++i) {
    // Find next vertex that not in set, which has smallest distance to existing
    // set of vertices.
    int min = INT_MAX; int minIdx = 0;
    for (int ii = 0; ii < N; ++ii) {
      if (!minset[ii] && dist[ii] < min) {
        minIdx = ii;
        min = dist[ii];
      }
    }

    // Add the new vertex to the set, and update all dist accordingly.
    minset[minIdx] = true;
    for (int j = 0; j < N; ++j) {
      if (grid[minIdx][j] && !minset[j] && grid[minIdx][j] < dist[j]) {
        dist[j] = grid[minIdx][j];
        parent[j] = minIdx;
      }
    }
  }

  int ret = 0;
  for (int i = 1; i < N; ++i) {
    ret += grid[parent[i]][i];
  }

  fout << ret << endl;
  cout << ret << endl;
}