Submission #2477445

---

Source Code Expand

#include<bits/stdc++.h>
using namespace std;
using ll = long long;

const int N = 2e5;
vector<int> c(N);
vector<int> used(N);
vector<int> gg[N];
int n, m, k;
void dfs(int i, int &un) {
  used[i] = 1;
  if(i < n && c[i] == 0) un++;
  for(int j : gg[i]) if(!used[j]) {
    dfs(j, un);
  }
}

/// --- Union Find Library {{{ ///

struct UF {
  int n;
  vector<int> par, rank;
  UF(int n): n(n), par(n, -1), rank(n, 0) {}
  int find(int x) {
    return par[x] < 0 ? x : par[x] = find(par[x]);
  }
  int size(int x) {
    return -par[find(x)];
  }
  bool same(int a, int b) {
    return find(a) == find(b);
  }
  void unite(int a, int b) {
    a = find(a);
    b = find(b);
    if(a == b) return;
    if(rank[a] > rank[b]) swap(a, b);
    par[b] += par[a];
    par[a] = b;
    if(rank[a] == rank[b]) rank[b]++;
  }
};

/// }}}--- ///

vector<tuple<int, int, int>> g;
ll solve() {
  ios::sync_with_stdio(false), cin.tie(0);
  cin >> n >> m >> k;
  set<int> cols;
  for(int i = 0; i < n; i++) {
    cin >> c[i];
    if(c[i] != 0) cols.emplace(c[i]);
  }
  for(int i = 0; i < m; i++) {
    int a, b, w;
    cin >> a >> b >> w;
    a--; b--;
    gg[a].emplace_back(b);
    gg[b].emplace_back(a);
    g.emplace_back(w, a, b);
    if(c[a]) g.emplace_back(0, a, n + c[a]);
    if(c[b]) g.emplace_back(0, b, n + c[b]);
  }

  int forests = 0;
  int no = 0;
  int fre = 0;
  for(int i = 0; i < n; i++) if(!used[i]) {
    forests++;
    int x = 0;
    dfs(i, x);
    if(x == 0) no++;
    fre += x;
    // cerr << "ref = " << fre << endl;
  }
  // cerr << forests << endl;
  // cerr << no << endl;
  //  << k - cols.size() << endl;
  if(no >= 2) return -1;
  if(fre < forests - 1 + (k - cols.size())) return -1;

  sort(begin(g), end(g));
  UF uf(n + k + 10);
  int cnt = 0;
  ll ans = 0;
  for(auto e : g) {
    int w = get<0>(e), a = get<1>(e), b = get<2>(e);
    if(uf.same(a, b)) continue;
    uf.unite(a, b);
    if(w) cnt++, ans += w;
    if(cnt == k - 1) break;
  }
  return ans;
}

int main() {
  cout << solve() << endl;
}

Submission Info

Judge Result