[CH6201]走廊泼水节 - 最小生成树 - Kruskal算法

【题目描述】
给定一棵N个节点的树，要求增加若干条边，把这棵树扩充为完全图，并满足图的唯一最小生成树仍然是这棵树。
求增加的边的权值总和最小是多少。
注意： 树中的所有边权均为整数，且新加的所有边权也必须为整数。
1≤N≤6000
【分析】
kruskal算法
按照边权从小到大排序给定的n-1条边，执行合并操作。
对于边权为w的边<u,v>，如果合并时结点u，v分属两个集合（连通块），则必须进行加边操作，为了保证唯一最小生成树是原树，新增加的边的边权为w+1，集合中的所有点都要对应连边，$ans+=(sz[x]+sz[y]-1)\ast (w+1)$
【代码如下】

#include <bits/stdc++.h>
using namespace std;
const int maxn = 6006;
struct E
{
    int u,v,w;
    friend bool operator < (const E &x,const E &y)
    {
        return x.w < y.w;
    }
}edge[maxn];
int sz[maxn],fa[maxn],n;
long long ans;
void init()
{
    memset(edge,0,sizeof(edge));
    memset(fa,0,sizeof(fa));
    memset(sz,0,sizeof(sz));
    ans = 0;
}
void read()
{
    scanf("%d",&n);
    for(int i=1;i<n;i++)
        scanf("%d%d%d",&edge[i].u,&edge[i].v,&edge[i].w);
}
int find(int x)
{
    return x == fa[x]?x:fa[x] = find(fa[x]);
}
void work()
{
    sort(edge+1,edge+n);
    for(int i=1;i<=n;i++)
        fa[i] = i, sz[i] = 1;
    for(int i=1;i<n;i++)
    {
        register int x = find(edge[i].u),y = find(edge[i].v);
        if(x!=y)
        {
            fa[x] = y;
            ans+=1LL*(sz[x]*sz[y]-1)*(edge[i].w+1);
            sz[y]+=sz[x];
        }
    }
    printf("%lld\n",ans);
}
int main()
{
    int t;
    scanf("%d",&t);
    while(t--)
    {
        init();
        read();
        work();
    }
    return 0;
}