Problem:


Score : 700 points

Problem Statement

Ringo has an undirected graph G with N vertices numbered 1,2,...,N and M edges numbered 1,2,...,M.
Edge i connects Vertex a_{i} and b_{i} and has a length of w_i.

Now, he is in the middle of painting these N vertices in K colors numbered 1,2,...,K. Vertex i is already painted in Color c_i, except when c_i = 0, in which case Vertex i is not yet painted.

After he paints each vertex that is not yet painted in one of the K colors, he will give G to Snuke.

Based on G, Snuke will make another undirected graph G' with K vertices numbered 1,2,...,K and M edges.
Initially, there is no edge in G'. The i-th edge will be added as follows:

Let x and y be the colors of the two vertices connected by Edge i in G.

Add an edge of length w_i connecting Vertex x and y in G'.

What is the minimum possible sum of the lengths of the edges in the minimum spanning tree of G'? If G' will not be connected regardless of how Ringo paints the vertices, print -1.

Constraints

1 \leq N,M \leq 10^{5}

1 \leq K \leq N

0 \leq c_i \leq K

1 \leq a_i,b_i \leq N

1 \leq w_i \leq 10^{9}

The given graph may NOT be simple or connected.

All input values are integers.

Partial Scores

In the test set worth 100 points, N = K and c_i = i.

In the test set worth another 100 points, c_i \neq 0.

In the test set worth another 200 points, c_i = 0.

Input

Input is given from Standard Input in the following format:

N M K
c_1 c_2 ... c_{N}
a_1 b_1 w_1
:
a_M b_M w_M

Output

Print the answer.

Sample Input 1

4 3 3
1 0 1 2
1 2 10
2 3 20
2 4 50

Sample Output 1

60

G' will only be connected when Vertex 2 is painted in Color 3.

Sample Input 2

5 2 4
0 0 0 0 0
1 2 10
2 3 10

Sample Output 2

-1

Regardless of how Ringo paints the vertices, two edges is not enough to connect four vertices as one.

Sample Input 3

9 12 9
1 2 3 4 5 6 7 8 9
6 9 9
8 9 6
6 7 85
9 5 545631016
2 1 321545
1 6 33562944
7 3 84946329
9 7 15926167
4 7 53386480
5 8 70476
4 6 4549
4 8 8

Sample Output 3

118901402

Sample Input 4

18 37 12
5 0 4 10 8 7 2 10 6 0 9 12 12 11 11 11 0 1
17 1 1
11 16 7575
11 15 9
10 10 289938980
5 10 17376
18 4 1866625
8 11 959154208
18 13 200
16 13 2
2 7 982223
12 12 9331
13 12 8861390
14 13 743
2 10 162440
2 4 981849
7 9 1
14 17 2800
2 7 7225452
3 7 85
5 17 4
2 13 1
10 3 45
1 15 973
14 7 56553306
16 17 70476
7 18 9
9 13 27911
18 14 7788322
11 11 8925
9 13 654295
2 10 9
10 1 545631016
3 4 5
17 12 1929
2 11 57
1 5 4
1 17 7807368

Sample Output 4

171



Sample input:
4 3 3
1 0 1 2
1 2 10
2 3 20
2 4 50


Sample output:
60

Submission status: Wrong Answer

Solution code:
class Unionfindtree:
    def __init__(self, number):
        self.par = [i for i in range(number)]
        self.rank = [0] * (number)

    def find(self, x):  # 親を探す
        if self.par[x] == x:
            return x
        else:
            self.par[x] = self.find(self.par[x])
            return self.par[x]

    def union(self, x, y):  # x,yを繋げる
        px = self.find(x)
        py = self.find(y)
        if px == py:
            return
        if self.rank[px] < self.rank[py]:
            self.par[px] = py
        else:
            self.par[py] = px
        if self.rank[px] == self.rank[py]:
            self.rank[px] += 1

    def connect(self, x, y):  # 親が同じかみる
        return self.find(x) == self.find(y)


N, M, K = map(int, input().split())
C = [int(i) for i in input().split()]
table = []
t = K
dd = {}
for i in range(M):
    s, u, c = map(int, input().split())
    a, b = C[s - 1], C[u - 1]
    if a == 0:
        a = dd.get(s, -1)
        if a == -1:
            t += 1
            dd[s] = t
            a = t
    elif b == 0:
        b = dd.get(u, -1)
        if b == -1:
            t += 1
            dd[u] = t
            b = t
    table.append((c, a, b))

tree = Unionfindtree(t+1)
table.sort()
ct = K - 1
ans = 0
for c, a, b in table:
    if not tree.connect(a, b) and ct > 0:
        tree.union(a, b)
        ans += c
        ct -= 1
if ct > 0:
    print(-1)
else:
    print(ans)

Submission status: Wrong Answer

Solution code:
N, M, K = map(int, input().split())
*C, = map(int, input().split())
G = [[] for i in range(N)]
E = []
for i in range(M):
    a, b, w = map(int, input().split())
    G[a-1].append((b-1, w))
    G[b-1].append((a-1, w))
    E.append((w, a-1, b-1))

def root(parent, x):
    if x == parent[x]:
        return x
    parent[x] = y = root(parent, parent[x])
    return y
def unite(parent, x, y):
    px = root(parent, x); py = root(parent, y)
    if px < py:
        parent[py] = px
    else:
        parent[px] = py

E.sort()
*p, = range(N)
res = 0
for w, a, b in E:
    if root(p, a) != root(p, b):
        unite(p, a, b)
        res += w
if N == K and all(C[i] == i+1 for i in range(N)):
    if not all(root(p, i) == 0 for i in range(N)):
        print(-1)
        exit(0)
    print(res)
else:
    if all(C[i] == 0 for i in range(N)):
        #if not all(root(p, i) == 0 for i in range(N)):
        #    print(-1)
        #    exit(0)
        print(0)
    if all(C[i] != 0 for i in range(N)):
        *p, = range(K)

        res = 0
        for w, a, b in E:
            ga = C[a]-1; gb = C[b]-1
            if root(p, ga) != root(p, gb):
                unite(p, ga, gb)
                res += w
        if not all(root(p, i) == 0 for i in range(K)):
            print(-1)
        else:
            print(res)

Submission status: Wrong Answer

Solution code:
N, M, K = map(int, input().split())
*C, = map(int, input().split())
G = [[] for i in range(N)]
E = []
for i in range(M):
    a, b, w = map(int, input().split())
    G[a-1].append((b-1, w))
    G[b-1].append((a-1, w))
    E.append((w, a-1, b-1))

def root(parent, x):
    if x == parent[x]:
        return x
    parent[x] = y = root(parent, parent[x])
    return y
def unite(parent, x, y):
    px = root(parent, x); py = root(parent, y)
    if px < py:
        parent[py] = px
    else:
        parent[px] = py

E.sort()
*p, = range(N)
res = 0
for w, a, b in E:
    if root(p, a) != root(p, b):
        unite(p, a, b)
        res += w
if N == K and all(C[i] == i+1 for i in range(N)):
    print(res)
else:
    if all(C[i] == 0 for i in range(N)):
        #if not all(root(p, i) == 0 for i in range(N)):
        #    print(-1)
        #    exit(0)
        print(0)
        exit(0)
    if all(C[i] != 0 for i in range(N)):
        *p, = range(K)

        res = 0
        for w, a, b in E:
            ga = C[a]-1; gb = C[b]-1
            if root(p, ga) != root(p, gb):
                unite(p, ga, gb)
                res += w
        if not all(root(p, i) == 0 for i in range(K)):
            print(-1)
        else:
            print(res)

Submission status: Accepted

Solution code:
class Unionfindtree:
    def __init__(self, number):
        self.par = [i for i in range(number)]
        self.rank = [0] * (number)

    def find(self, x):  # 親を探す
        if self.par[x] == x:
            return x
        else:
            self.par[x] = self.find(self.par[x])
            return self.par[x]

    def union(self, x, y):  # x,yを繋げる
        px = self.find(x)
        py = self.find(y)
        if px == py:
            return
        if self.rank[px] < self.rank[py]:
            self.par[px] = py
        else:
            self.par[py] = px
        if self.rank[px] == self.rank[py]:
            self.rank[px] += 1

    def connect(self, x, y):  # 親が同じかみる
        return self.find(x) == self.find(y)


N, M, K = map(int, input().split())
C = [int(i) for i in input().split()]
table = []
t = K
for i in range(N):
    if C[i]==0:
        t+=1
        C[i]=t
for i in range(M):
    s, u, c = map(int, input().split())
    a, b = C[s - 1], C[u - 1]
    table.append((c, a, b))

tree = Unionfindtree(t+1)
table.sort()
ct = K - 1
ans = 0
for c, a, b in table:
    if not tree.connect(a, b) and ct > 0:
        tree.union(a, b)
        ans += c
        ct -= 1
if ct > 0:
    print(-1)
else:
    print(ans)

Submission status: Runtime Error

Solution code:
#include "bits/stdc++.h"

using namespace std;
#ifdef _DEBUG
#define _GLIBCXX_DEBUG
#include "dump.hpp"
#else
#define dump(...)
#endif

#define int long long
#define ll long long
#define ll1 1ll
#define ONE 1ll
#define DBG 1
#define rep(i, a, b) for (int i = (a); i < (b); i++)
#define rrep(i, a, b) for (int i = (b)-1; i >= (a); i--)
#define loop(n) rep(loop, (0), (n))
#define all(c) begin(c), end(c)
const int INF =
sizeof(int) == sizeof(long long) ? 0x3f3f3f3f3f3f3f3fLL : 0x3f3f3f3f;
const int MOD = (int)(1e9) + 7;
const double PI = acos(-1);
const double EPS = 1e-9;
#define fi first
#define se second
#define pb push_back
#define eb emplace_back
using pii = pair<int, int>;
// template<class T> ostream &operator<<(ostream &os,T &t){dump(t);return os;}
template <typename T, typename S>
istream &operator>>(istream &is, pair<T, S> &p) {
	is >> p.first >> p.second;
	return is;
}
template <typename T, typename S>
ostream &operator<<(ostream &os, pair<T, S> &p) {
	os << p.first << " " << p.second;
	return os;
}

template <typename T> void printvv(const vector<vector<T>> &v) {
	cerr << endl;
	rep(i, 0, v.size()) rep(j, 0, v[i].size()) {
		if (typeid(v[i][j]).name() == typeid(INF).name() and v[i][j] == INF) {
			cerr << "INF";
		}
		else
			cerr << v[i][j];
		cerr << (j == v[i].size() - 1 ? '\n' : ' ');
	}
	cerr << endl;
}
/*
typedef __int128_t Int;
std::ostream &operator<<(std::ostream &dest, __int128_t value) {
std::ostream::sentry s(dest);
  if (s) {
	__uint128_t tmp = value < 0 ? -value : value;
	char buffer[128];
	char *d = std::end(buffer);
	do {
	  --d;
	  *d = "0123456789"[tmp % 10];
	  tmp /= 10;
	} while (tmp != 0);
	if (value < 0) {
	  --d;
	  *d = '-';
	}
	int len = std::end(buffer) - d;
	if (dest.rdbuf()->sputn(d, len) != len) {
	  dest.setstate(std::ios_base::badbit);
	}
  }
  return dest;
}

__int128 parse(string &s) {
  __int128 ret = 0;
  for (int i = 0; i < s.length(); i++)
	if ('0' <= s[i] && s[i] <= '9')
	  ret = 10 * ret + s[i] - '0';
  return ret;
}
*/

#ifndef _DEBUG
#define printvv(...)
#endif
void YES(bool f) { cout << (f ? "YES" : "NO") << endl; }
void Yes(bool f) { cout << (f ? "Yes" : "No") << endl; }
template <class T> bool chmax(T &a, const T &b) {
	if (a < b) {
		a = b;
		return true;
	}
	return false;
}
template <class T> bool chmin(T &a, const T &b) {
	if (a > b) {
		a = b;
		return true;
	}
	return false;
}
using Weight = int;
using Flow = int;
struct Edge {
	int s, d; Weight w; Flow c;
	Edge() {};
	Edge(int s, int d, Weight w = 1) : s(s), d(d), w(w), c(w) {};
};
bool operator<(const Edge &e1, const Edge &e2) { return e1.w < e2.w; }
bool operator>(const Edge &e1, const Edge &e2) { return e2 < e1; }
inline ostream &operator<<(ostream &os, const Edge &e) { return (os << '(' << e.s << ", " << e.d << ", " << e.w << ')'); }

using Edges = vector<Edge>;
using Graph = vector<Edges>;
using Array = vector<Weight>;
using Matrix = vector<Array>;

void addArc(Graph &g, int s, int d, Weight w = 1) {
	g[s].emplace_back(s, d, w);
}
void addEdge(Graph &g, int a, int b, Weight w = 1) {
	addArc(g, a, b, w);
	addArc(g, b, a, w);
}
struct DisjointSet {
	// ２つのグループを１つにまとめる　と　2つの要素が同じグループに所属しているかどうかを判定する
	vector<int> rank, p, S; // p->parent  S[findSet(v)] ->連結成分の大きさ
	DisjointSet() {}
	DisjointSet(int size) {
		S.resize(size, 1);
		rank.resize(size, 0);
		p.resize(size, 0);
		rep(i, 0, size) makeSet(i);
	}
	void makeSet(int x) {
		p[x] = x;
		rank[x] = 0;
	}
	bool same(int x, int y) { // 判定する
		return findSet(x) == findSet(y);
	}
	void unite(int x, int y) { // 連結するときにはこれを使う
		if (same(x, y))
			return;
		link(findSet(x), findSet(y));
	}
	void link(int x, int y) {
		if (rank[x] > rank[y]) {
			p[y] = x;
		}
		else {
			p[x] = y;
			if (rank[x] == rank[y]) {
				rank[y]++;
			}
		}

		S[x] = S[y] = S[x] + S[y];
	}
	int findSet(int x) {
		if (x != p[x]) {
			p[x] = findSet(p[x]); // path compression
		}
		return p[x];
	}
	int connectedComponentSize(int x) { return S[findSet(x)]; }
};

int kruskal(int N,int K, vector<Edge> &edges) {
	int totalCost = 0;
	sort(all(edges));
	DisjointSet dset(N);
	int k = 0;
	rep(i, 0, edges.size()) {
		Edge e = edges[i];
		if (!dset.same(e.s, e.d)) {
			totalCost += e.w;
			dset.unite(e.s, e.d);
			k++;
			if (k == K - 1)break;
		}
	}
	if (k != K - 1)totalCost = -1;
	return totalCost;
}

signed main(signed argc, char *argv[]) {
	cin.tie(0);
	ios::sync_with_stdio(false);
	cout << fixed << setprecision(12);

	int N, M; cin >> N >> M;
	int K; cin >> K;
	vector<int>c(N);
	int x = K;
	rep(i, 0, N) {
		cin >> c[i];
		c[i]--;
		if (c[i] == -1)c[i] = x++;
	}

	Edges es;
	loop(M) {
		int a, b, w; cin >> a >> b >> w;
		a--, b--;
		es.eb(c[a], c[b], w);
	}
	cout << kruskal(x,K, es) << endl;
	return 0;
}

Submission status: Runtime Error

Solution code:
N, M, K = map(int, input().split())
*C, = map(int, input().split())
G = [[] for i in range(N)]
E = []
for i in range(M):
    a, b, w = map(int, input().split())
    G[a-1].append((b-1, w))
    G[b-1].append((a-1, w))
    E.append((w, a-1, b-1))

def root(parent, x):
    if x == parent[x]:
        return x
    parent[x] = y = root(parent, parent[x])
    return y
def unite(parent, x, y):
    px = root(parent, x); py = root(parent, y)
    if px < py:
        parent[py] = px
    else:
        parent[px] = py

E.sort()
*p, = range(N)
res = 0
for w, a, b in E:
    if root(p, a) != root(p, b):
        unite(p, a, b)
        res += w
if N == K and all(C[i] == i+1 for i in range(N)):
    if not all(root(p, i) == 0 for i in range(N)):
        print(-1)
        exit(0)
    print(res)
else:
    if all(C[i] == 0 for i in range(N)):
        count = 0
        *p, = range(N)
        res = 0
        for w, a, b in E:
            if root(p, a) != root(p, b) and count < K:
                unite(p, a, b)
                res += w
                count += 1
        if count == K:
            print(res)
        else:
            print(-1)
    elif all(C[i] != 0 for i in range(N)):
        *p, = range(K)

        res = 0
        for w, a, b in E:
            ga = C[a]-1; gb = C[b]-1
            if root(p, ga) != root(p, gb):
                unite(p, ga, gb)
                res += w
        if not all(root(p, i) == 0 for i in range(K)):
            print(-1)
        else:
            print(res)
    else:
        exit(1)

Submission status: Runtime Error

Solution code:
N, M, K = map(int, input().split())
*C, = map(int, input().split())
G = [[] for i in range(N)]
E = []
for i in range(M):
    a, b, w = map(int, input().split())
    G[a-1].append((b-1, w))
    G[b-1].append((a-1, w))
    E.append((w, a-1, b-1))

def root(parent, x):
    if x == parent[x]:
        return x
    parent[x] = y = root(parent, parent[x])
    return y
def unite(parent, x, y):
    px = root(parent, x); py = root(parent, y)
    if px < py:
        parent[py] = px
    else:
        parent[px] = py

E.sort()
*p, = range(N)
res = 0
for w, a, b in E:
    if root(p, a) != root(p, b):
        unite(p, a, b)
        res += w
if N == K and all(C[i] == i+1 for i in range(N)):
    if not all(root(p, i) == 0 for i in range(N)):
        print(-1)
        exit(0)
    print(res)
else:
    if all(C[i] == 0 for i in range(N)):
        count = 0
        *p, = range(N)
        res = 0
        for w, a, b in E:
            if root(p, a) != root(p, b) and count < K:
                unite(p, a, b)
                res += w
                count += 1
        print(res)
    elif all(C[i] != 0 for i in range(N)):
        *p, = range(K)

        res = 0
        for w, a, b in E:
            ga = C[a]-1; gb = C[b]-1
            if root(p, ga) != root(p, gb):
                unite(p, ga, gb)
                res += w
        if not all(root(p, i) == 0 for i in range(K)):
            print(-1)
        else:
            print(res)
    else:
        exit(1)