Kruskal算法 走廊潑水節 CH6201

問題描述
給定一棵N個節點的樹，要求增加若干條邊，把這棵樹擴充爲完全圖，
並滿足圖的唯一最小生成樹仍然是這棵樹。
求增加的邊的權值總和最小是多少。
注意： 樹中的所有邊權均爲整數，且新加的所有邊權也必須爲整數。

輸入格式
第一行包含整數t，表示共有t組測試數據。

對於每組測試數據，第一行包含整數N。

接下來N-1行，每行三個整數X,Y,Z，表示X節點與Y節點之間存在一條邊，長度爲Z。

輸出格式
每組數據輸出一個整數，表示權值總和最小值。

每個結果佔一行。

數據範圍
1 ≤ N ≤ 6000
1 ≤ Z ≤ 100

輸入樣例：
2
3
1 2 2
1 3 3
4
1 2 3
2 3 4
3 4 5 
輸出樣例：
4
17 

思路
將每條邊安從小到大排序，將這些邊按照Kruskal算法做最小生成樹的操作，保證每個集合內的點都是完全圖，所以在合併集合的同時計算連接所有的點所需要的邊權之和，連接兩個集合的邊數爲兩個集合的點數相乘減去一，邊權爲當前的最小生成樹的這個邊的權值加一（滿足整數要求）

#include <iostream>
#include <cstring>
#include <algorithm>
#include <cstdio>

using namespace std;

typedef long long ll;
const int N = 6010;

struct node
{
	int x, y, z;
	bool operator<(const node &a) const
	{
		return z < a.z;
	}
} e[N];

int f[N], s[N];
ll ans;
int n;

void init()
{
	for (int i = 0; i <= N - 1; i++)
	{
		s[i] = 1;
		f[i] = i;
	}
}

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

int main()
{
	int T;
	scanf("%d", &T);
	while (T--)
	{
		init();
		ans = 0;
		scanf("%d", &n);

		for (int i = 0; i < n - 1; i++)
		{
			int a, b, c;
			scanf("%d%d%d", &a, &b, &c);
			e[i] = {a, b, c};
		}

		sort(e, e + n - 1);

		for (int i = 0; i < n - 1; i++)
		{
			int x = e[i].x;
			int y = e[i].y;
			int z = e[i].z;
			x = find(x);
			y = find(y);
			if (x != y)
			{
				ans += (1ll * z + 1) * (s[x] * s[y] - 1);
				f[x] = f[y];
				s[y] += s[x];
			}
		}
		printf("%d\n", ans);
	}
	return 0;
}