Problem p03517 (C++)

Buggy submission (status: Runtime Error):
/*
 _   ________   ___       __ ___   _           _  _
| | / /|  _  \ / _ \     / // _ \ | |         | |(_)
| |/ / | | | |/ /_\ \   / // /_\ \| | __ __ _ | | _
|    \ | | | ||  _  |  / / |  _  || |/ // _` || || |
| |\  \| |/ / | | | | / /  | | | ||   <| (_| || || |
\_| \_/|___/  \_| |_//_/   \_| |_/|_|\_\\__,_||_||_|
*/
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
using namespace __gnu_pbds;
using namespace std;
#define int long long
const int N = 1e5 + 50;
const int mod = 1e9 + 7;
const int in = 1e9 + 9;
#define size fksjfs
typedef pair<int, int> ii;
int n, m, k, size[N], col[N];
int parents[N];
pair<ii, int> edge[N];
#define fi first
#define se second
void Init(int n) {
  for (int i = 1; i <= n; i++)
    parents[i] = i, size[i] = 1;
}
int Find(int x) {
  return parents[x] == x ? x : (parents[x] = Find(parents[x]));
}
void Union(int i, int j) {
  int x = Find(i);
  int y = Find(j);
  if (x == y)
    return;
  if (size[x] >= size[y]) {
    parents[y] = x;
    size[x] += size[y];
  } else {
    parents[x] = y;
    size[y] += size[x];
  }
}
signed main() {
  cin >> n >> m >> k;
  for (int i = 1; i <= n; i++)
    cin >> col[i];
  for (int i = 1; i <= m; i++) {
    int u, v, w;
    cin >> u >> v >> w;
    if (col[u])
      u = col[u];
    else
      u += k;
    if (col[v])
      v = col[v];
    else
      v += k;
    edge[i] = {{u, v}, w};
  }
  Init(n + k);
  sort(edge + 1, edge + m + 1,
       [](pair<ii, int> _a, pair<ii, int> _b) { return _a.se < _b.se; });
  int ans = 0;
  int cnt = 0;
  for (int i = 1; i <= m; i++) {
    if (cnt == k - 1)
      continue;
    int u = edge[i].fi.fi, v = edge[i].fi.se;
    if (Find(u) != Find(v)) {
      cnt++;
      Union(u, v);
      ans += edge[i].se;
    }
  }
  if (cnt != k - 1)
    return cout << -1, 0;
  cout << ans;
  return 0;
}


Fixed submission:
/*
 _   ________   ___       __ ___   _           _  _
| | / /|  _  \ / _ \     / // _ \ | |         | |(_)
| |/ / | | | |/ /_\ \   / // /_\ \| | __ __ _ | | _
|    \ | | | ||  _  |  / / |  _  || |/ // _` || || |
| |\  \| |/ / | | | | / /  | | | ||   <| (_| || || |
\_| \_/|___/  \_| |_//_/   \_| |_/|_|\_\\__,_||_||_|
*/
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
using namespace __gnu_pbds;
using namespace std;
#define int long long
const int N = 5e5 + 50;
const int mod = 1e9 + 7;
const int in = 1e9 + 9;
#define size fksjfs
typedef pair<int, int> ii;
int n, m, k, size[N], col[N];
int parents[N];
pair<ii, int> edge[N];
#define fi first
#define se second
void Init(int n) {
  for (int i = 1; i <= n; i++)
    parents[i] = i, size[i] = 1;
}
int Find(int x) {
  return parents[x] == x ? x : (parents[x] = Find(parents[x]));
}
void Union(int i, int j) {
  int x = Find(i);
  int y = Find(j);
  if (x == y)
    return;
  if (size[x] >= size[y]) {
    parents[y] = x;
    size[x] += size[y];
  } else {
    parents[x] = y;
    size[y] += size[x];
  }
}
signed main() {
  cin >> n >> m >> k;
  for (int i = 1; i <= n; i++)
    cin >> col[i];
  for (int i = 1; i <= m; i++) {
    int u, v, w;
    cin >> u >> v >> w;
    if (col[u])
      u = col[u];
    else
      u += k;
    if (col[v])
      v = col[v];
    else
      v += k;
    edge[i] = {{u, v}, w};
  }
  Init(n + k);
  sort(edge + 1, edge + m + 1,
       [](pair<ii, int> _a, pair<ii, int> _b) { return _a.se < _b.se; });
  int ans = 0;
  int cnt = 0;
  for (int i = 1; i <= m; i++) {
    if (cnt == k - 1)
      continue;
    int u = edge[i].fi.fi, v = edge[i].fi.se;
    if (Find(u) != Find(v)) {
      cnt++;
      Union(u, v);
      ans += edge[i].se;
    }
  }
  if (cnt != k - 1)
    return cout << -1, 0;
  cout << ans;
  return 0;
}