Codeforces Round #441F: Royal Questions

http://codeforces.com/contest/875/problem/F

princeを点、princessの二人のprinceの候補を辺とみます。それからkruskalっぽいことをします。具体的には
1.違う連結成分どおしで、二つともlockedじゃなければ繋ぐ。
2.同じ連結成分どおしで、lockedじゃなければ繋いでからlockedにする。
この貪欲が成り立つ理由は、連結成分内ではkruskalの要領で最大を示せ、違う連結成分をまたぐ辺が使われないことは、2つの連結成分内にその辺よりもコストが低い辺は存在しないためです。

わかるんだけどもこういうのがコンテスト中に通せるかというと微妙。通せるようになりましょう。

int N, M, E;
edge es[MAX_N];

int kruskal() {
	sort(es, es + M, comp);
	init(N);
	int res = 0;
	for(int i = 0; i < M; i++) {
		edge e = es[i];
		int x = find(e.u), y = find(e.v);
		if(!locked[x] && !locked[y]) {
			if(x != y) {
				unite(x, y);
				res += e.cost;
			}
			else {
				locked[x] = true;
				res += e.cost;
			}
		}
		else if(!locked[x]) {
			unite(x, y);
			res += e.cost;
			locked[find(x)] = true;
		}
		else if(!locked[y]) {
			unite(x, y);
			res += e.cost;
			locked[find(y)] = true;
		}
	}
	return res;
}


void add_edge(int s, int t, int cost) {
	es[E++] = edge{s, t, cost};
}

void solve() {
	E = 0;
	cin >> N >> M;
	rep(i, 0, M) {
		int a, b, c; cin >> a >> b >> c;
		a--; b--;
		add_edge(a, b, c);
	}
	cout << kruskal() << "\n";
}