最小生成树练习

走廊泼水节
  题意: 给定一棵N个节点的树，要求增加若干边, 使其扩充为完全图, 并且图的唯一最小生成树仍然是这棵树, 求增加边的权值总和是多少.
  分析: 考虑一个类似$krusal$的过程, 先对边权进行排序, 依次扫描每条边, 对$(x,y,z)$x所在的并查集为$S_x$, y所在并查集为$S_y$, 合并$S_x$和$S_y$, 对于两个集合的路径, 两个集合间一共会增加$|S_x\ast S_y|-1$条边, 所以答案增加$(z+1)\ast (|S_x|\ast |S_y|-1)$

#include <bits/stdc++.h>
using namespace std;
#define ll long long
#define inc(i,l,r) for (int i = l; i <= r; ++i)
const int N = 6005;
template <typename T>
void read(T &x)
{
    x = 0; ll f = 1; char ch = getchar();
    while (!isdigit(ch)) { if (ch == '-') f = -1; ch = getchar(); }
    while (isdigit(ch)) x = (x<<3) + (x<<1) + ch - '0', ch = getchar();
    x *= f;
}

struct edge {
    int u, v, w;
    friend bool operator < (edge a, edge b) {
        return a.w < b.w;
    }
}e[N];

int t, n;
int f[N], s[N];

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

int solve() {
    sort(e + 1, e + n);
    int ans = 0;
    for (int i = 1; i <= n; ++i) f[i] = i, s[i] = 1;
    for (int i = 1; i < n; ++i) {
        int x = Find(e[i].u);
        int y = Find(e[i].v);
        if (x != y) {
            ans += (e[i].w + 1) * (s[x] * s[y] - 1);
            s[y] += s[x];
            f[x] = y;
        }
    }
    return ans;
}

int main()
{
    cin >> t;
    while (t--) {
        cin >> n;
        for (int i = 1; i < n; ++i)
            scanf("%d%d%d", &e[i].u, &e[i].v, &e[i].w);
        cout << solve() << endl;
    }
    return 0;
}