Problem p03517 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <bitset>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <deque>
#include <fstream>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>
using namespace std;

#define REP(i, n) for (int i = 0; i < n; ++i)
#define FOR(i, a, b) for (int i = a; i <= b; ++i)
#define FORR(i, a, b) for (int i = a; i >= b; --i)
#define ALL(c) (c).begin(), (c).end()

typedef long long ll;
typedef vector<int> VI;
typedef vector<ll> VL;
typedef vector<VL> VVL;
typedef vector<VI> VVI;
typedef pair<int, int> P;
typedef pair<ll, ll> PL;

struct union_find {
  VI par;

  void init(int n) {
    par.resize(n);
    REP(i, n) par[i] = i;
  }

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

  bool same(int x, int y) { return find(x) == find(y); }

  void merge(int x, int y) {
    x = find(x);
    y = find(y);
    if (x == y)
      return;
    par[x] = y;
  }
};

int main() {
  int n, m, k;
  cin >> n >> m >> k;

  VI c(n);
  REP(i, n) {
    scanf("%d", &c[i]);
    c[i]--;
  }

  VI a(m), b(m);
  VL w(m);
  REP(i, m) {
    scanf("%d %d %lld", &a[i], &b[i], &w[i]);
    a[i]--;
    b[i]--;
  }

  VI emp(n, -1);
  int x = k;
  REP(i, n) {
    if (c[i] == -1) {
      emp[i] = x++;
    }
  }

  VI ind(m);
  REP(i, m) ind[i] = i;

  sort(ALL(ind), [&](int x, int y) { return w[x] < w[y]; });

  union_find uf;
  uf.init(n);

  ll ans = 0;
  x = 0;

  REP(i, m) {
    int p = (c[a[ind[i]]] == -1 ? emp[a[ind[i]]] : c[a[ind[i]]]);
    int q = (c[b[ind[i]]] == -1 ? emp[b[ind[i]]] : c[b[ind[i]]]);
    if (!uf.same(p, q)) {
      x++;
      ans += w[ind[i]];
      uf.merge(p, q);

      if (x == k - 1) {
        cout << ans << endl;
        return 0;
      }
    }
  }
  cout << -1 << endl;

  return 0;
}


Fixed submission:
#include <algorithm>
#include <bitset>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <deque>
#include <fstream>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>
using namespace std;

#define REP(i, n) for (int i = 0; i < n; ++i)
#define FOR(i, a, b) for (int i = a; i <= b; ++i)
#define FORR(i, a, b) for (int i = a; i >= b; --i)
#define ALL(c) (c).begin(), (c).end()

typedef long long ll;
typedef vector<int> VI;
typedef vector<ll> VL;
typedef vector<VL> VVL;
typedef vector<VI> VVI;
typedef pair<int, int> P;
typedef pair<ll, ll> PL;

struct union_find {
  VI par;

  void init(int n) {
    par.resize(n);
    REP(i, n) par[i] = i;
  }

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

  bool same(int x, int y) { return find(x) == find(y); }

  void merge(int x, int y) {
    x = find(x);
    y = find(y);
    if (x == y)
      return;
    par[x] = y;
  }
};

int main() {
  int n, m, k;
  cin >> n >> m >> k;

  VI c(n);
  REP(i, n) {
    scanf("%d", &c[i]);
    c[i]--;
  }

  VI a(m), b(m);
  VL w(m);
  REP(i, m) {
    scanf("%d %d %lld", &a[i], &b[i], &w[i]);
    a[i]--;
    b[i]--;
  }

  VI emp(n, -1);
  int x = k;
  REP(i, n) {
    if (c[i] == -1) {
      emp[i] = x++;
    }
  }

  VI ind(m);
  REP(i, m) ind[i] = i;

  sort(ALL(ind), [&](int x, int y) { return w[x] < w[y]; });

  union_find uf;
  uf.init(x);

  ll ans = 0;
  x = 0;

  REP(i, m) {
    int p = (c[a[ind[i]]] == -1 ? emp[a[ind[i]]] : c[a[ind[i]]]);
    int q = (c[b[ind[i]]] == -1 ? emp[b[ind[i]]] : c[b[ind[i]]]);
    if (!uf.same(p, q)) {
      x++;
      ans += w[ind[i]];
      uf.merge(p, q);

      if (x == k - 1) {
        cout << ans << endl;
        return 0;
      }
    }
  }
  cout << -1 << endl;

  return 0;
}