走廊泼水节-Kruskal算法扩展

走廊泼水节

题目

346. 走廊泼水节

思路

​ 考虑\(N\)个节点数，将最小生成树扩充为完全图（\(N\)个点， 每个点\(N - 1\)条边）， 我们从Kruskal算法的思路开始思考, Kruskal算法通过不断的合并点集来建立最小生成树， 在合并的过程中，两个点集之间一定是没有任何边的，所有我们考虑在合并的过程中将图扩展为完全图， 考虑权值最小， 则取\(w(u, v) + 1\), 边数则为\(size(u) \times size(v) - 1\)， 由于Kruskal算法是正确的，并且保证所有点集合都会被合并到一个集合，每次合并保证合并完的集合是一个完全图， 所以会得到完全图。

思路

#include <bits/stdc++.h>

using i64 = long long;

const int N = 6e3 + 10;

struct Edge {
	int a, b, w;

	void input() {
		std::cin >> a >> b >> w;
	}

	void out() {
		std::cout << a << " " << b << " " << w << "\n";
	}

	bool operator<(const Edge &t) const {
		return w < t.w;
	}
}edge[N];


int n, m;

int p[N], sz[N];

int find(int x) {
	if (x != p[x]) p[x] = find(p[x]);
	return p[x];
}

i64 kruskal() {
	std::sort(edge + 1, edge + n);	

	i64 res = 0;

	for (int i = 1; i <= n - 1; i ++) {
		int a = edge[i].a, b = edge[i].b, w = edge[i].w;

		// edge[i].out();
		
		int fa = find(a), fb = find(b);

		if (fa != fb) {
			res += (sz[fa] * sz[fb] - 1) * (w + 1);
			p[fa] = fb;
			sz[fb] += sz[fa];
		}
	}

	return res;
}

void solve() {
	std::cin >> n;

	for (int i = 1; i <= n; i ++) {		
		p[i] = i, sz[i] = 1;
	}

	for (int i = 1; i <= n - 1; i ++) {
		edge[i].input();
	}

	std::cout << kruskal() << "\n";
}

int main() {
	int _;

	std::cin >> _;

	while (_ --) {
		solve();
	}
}