Submission #3527316

---

Source Code Expand

#include <bits/stdc++.h>
using namespace std;
using lint = long long int;
template<class T = int> using V = vector<T>;
template<class T = int> using VV = V< V<T> >;
template<class T> void assign(V<T>& v, int n, const T& a = T()) { v.assign(n, a); }
template<class T, class... Args> void assign(V<T>& v, int n, const Args&... args) { v.resize(n); for (auto&& e : v) assign(e, args...); }


struct QU {
  V<> par, rank, _size;

  QU(int n) : par(n), rank(n), _size(n, 1) { 
    iota(par.begin(), par.end(), 0);
  }

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

  bool same(int a, int b) { return find(a) == find(b); }

  int size(int a) { return _size[find(a)]; }

  void unite(int a, int b) {
    a = find(a), b = find(b);
    if (a == b) return;
    if (rank[a] < rank[b]) {
      par[a] = b;
      _size[b] += _size[a];
    } else {
      par[b] = a;
      _size[a] += _size[b];
    }
    if (rank[a] == rank[b]) rank[a]++;
  }
};

int main() {
  cin.tie(nullptr); ios_base::sync_with_stdio(false);
  int n, m, k; cin >> n >> m >> k;
  V<> c(n); for (int i = 0; i < n; ++i) cin >> c[i], --c[i];
  struct Edge { int u, v, w; };
  V<Edge> es;
  QU qu(n);
  for (int i = 0; i < m; ++i) {
    int u, v, w; cin >> u >> v >> w, --u, --v;
    es.push_back({u, v, w});
    qu.unite(u, v);
  }
  sort(begin(es), end(es), [](const Edge& lhs, const Edge& rhs) { return lhs.w < rhs.w; });
  V<> vs(n); iota(begin(vs), end(vs), 0);
  sort(begin(vs), end(vs), [&](int u, int v) { return c[u] < c[v]; });
  for (int i = 0; i < n - 1; ++i) if (c[vs[i]] != -1 and c[vs[i]] == c[vs[i + 1]]) qu.unite(vs[i], vs[i + 1]);
  set<int> sc;
  for (int i = 0; i < n; ++i) sc.insert(qu.find(i));
  int C = sc.size(), D = 0;
  for (int i = 0; i < n; ++i) D += c[i] == -1;
  set<int>().swap(sc);
  for (int i = 0; i < n; ++i) if (c[i] != -1) sc.insert(c[i]);
  int E = k - sc.size();
  if (C - 1 > D - E) return cout << -1 << '\n', 0;
  qu = QU(k + n);
  lint res = 0;
  int ce = 0;
  for (const auto& e : es) {
    int u = c[e.u] != -1 ? c[e.u] : k + e.u;
    int v = c[e.v] != -1 ? c[e.v] : k + e.v;
    if (qu.same(u, v)) continue;
    qu.unite(u, v);
    res += e.w;
    ++ce;
    if (ce == k - 1) break;
  }
  cout << res << '\n';
}

Submission Info

Judge Result