Vijos 1579 最小生成树kruskal的变形

问题链接：https://vijos.org/p/1579

大意为找到一颗最小生成树所对应的边权之和最小的完全图。
kruskal的算法，是每一次连通两个连通块，且边权尽量小，我们假设这个边权是L
那么 这两个连通块之间的其他连线肯定>L
那么每次我们连通两个连通块时，就可以得到这两个连通块其他线的最小长度L+1

#include <iostream>
#include <cstdio>
#include <cstdlib>
#include <algorithm>
using namespace std;
#define MAX 20001
typedef long long LL;   
struct node
{
    LL a;
    LL b;
    LL l;
}line[MAX];
bool operator<(node a,node b)
{
    return a.l < b.l;
}
LL n;
LL mot[MAX],num[MAX],ans;
LL find(LL x)
{
    if ( x != mot[x] )
    {
        mot[x] = find(mot[x]);
    }
    return mot[x];
}
void union(LL a,LL b,LL l)
{
    a = find(a);
    b = find(b);
    ans += (num[a] * num[b] -1) * (l+1);  //将两个要合并的并查集的除最小生成树中那条边之外的边权设置为l+1
    mot[b] = a;
    num[a] += num[b];
}
void init()
{
    cin>>n;
    for (LL i = 1; i < n; i++)
    {
        scanf("%I64d%I64d%I64d",&line[i].a,&line[i].b,&line[i].l);
    }
}
void work()
{
    for (LL i = 1; i <= n; i++)
    {
        mot[i] = i;       //初始化并查集
        num[i] = 1;       //维护并查集大小
    }
    sort(line+1,line+n);
    for (LL i = 1; i < n; i++)
    {//合并两个连通块时，增加的边为最小生成树中的一条边l和剩下的长为l+1的边
        ans += line[i].l;   
        union(line[i].a,line[i].b,line[i].l);
    }
}
void print()
{
    cout<<ans;
}
int main()
{
    init();
    work();
    print();
    return 0;
}