AcWing 346. 走廊泼水节(最小生成树)

Description

给定一棵 n 个节点的树，要求增加若干条边，把这棵树扩充为完全图，并满足图的唯一最小生成树仍然是这棵树。

求增加的边的权值总和最小是多少。

注意： 树中的所有边权均为整数，且新加的所有边权也必须为整数。

Input

第一行包含整数 T，表示共有 T 组测试数据。

对于每组测试数据，第一行包含整数 n ( 1 ≤ n ≤ 6000 ) 。

接下来 n - 1 行，每行三个整数 x, y, z，表示 x 节点与 y 节点之间存在一条边，长度为 z  ( 1 ≤ z ≤ 100 ) 。

Output

每组数据一行结果：一个整数，表示权值总和最小值

Sample Input

2
3
1 2 2
1 3 3
4
1 2 3
2 3 4
3 4 5

Sample Output

4
17

思路：设cnt[]存储各连通块的节点数量，在用kruskal算法求最小生成树时，当每次为最小生成树添加一条边权为w的边，说明这条边所连接的两个节点所在的连通块不连通，记这两个连通块为a,b，若想让这两个连通块成为一个完全图至少需要添加cnt[a]*cnt[b]-1条边(减去当前所添加边权为w的一条边)，若想让最小生成树不变，则新添加边的权重必须大于w，否则可以将w替换成新添加的其中一条边，则会获得一颗新的最小生成树。为了使得最小生成树不变且新添加边的权重最小，因此新添加边的权重为w+1。
因此我们只需要在kruskal算法每次为最小生成树新加一条权重为w的边时，计算当前两个连通a，b块成为完全图所需添加的边的最小权重即可，最小权重即为(cnt[a] * cnt[b] -1) * (w + 1)

代码如下：

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

#define N 6020

int p[N], cnt[N];
int find(int x) {
	if (x == p[x])return x;
	return p[x] = find(p[x]);
}

struct ver {
	int a, b, w;
	bool operator < (const ver& W)const {
		return w < W.w;
	}
}e[N];

int main() {
	int t, n, idx, ans;
	int a, b, c;
	cin >> t;
	while (t--) {
		idx = 0, ans = 0;
		cin >> n;
		for (int i = 1; i <= n; i++) { p[i] = i, cnt[i] = 1; }

		for (int i = 1; i < n; i++) {
			cin >> a >> b >> c;
			e[idx++] = { a,b,c };
		}
		sort(e, e + idx);
		for (int i = 0; i < idx; i++) {
			int a = e[i].a, b = e[i].b, w = e[i].w;
			int fa = find(a), fb = find(b);
			
			if (fa == fb)continue;		
			ans += ((cnt[fa] * cnt[fb] - 1) * (w + 1));
			p[fa] = fb;
			cnt[fb] += cnt[fa];
		}

		cout << ans << endl;
	}	
}