【ACWing】346. 走廊泼水节

题目地址：

https://www.acwing.com/problem/content/348/

给定一棵$N$个节点的树，要求增加若干条边，把这棵树扩充为完全图，并满足图的唯一最小生成树仍然是这棵树。求增加的边的权值总和最小是多少。注意：树中的所有边权均为整数，且新加的所有边权也必须为整数。

输入格式：
第一行包含整数$t$，表示共有$t$组测试数据。对于每组测试数据，第一行包含整数$N$。接下来$N-1$行，每行三个整数$X,Y,Z$，表示$X$节点与$Y$节点之间存在一条边，长度为$Z$。

输出格式：
每组数据输出一个整数，表示权值总和最小值。每个结果占一行。

数据范围：
$1\le N\le 6000$
$1\le Z\le 100$

其实就是每次加边的时候，考虑一下这个边与树形成的环中，除了这条边以外的最长边是哪条，那么新加的这条边的长度一定要恰好比最长哪条长$1$，这样才会导致树是唯一的最小生成树。可以换一种思路来求，我们可以想象Kruskal算法的最小生成树的构造，每次加一条边的时候就会连通两个连通块，那么要使得加的这条边恰好在原树里，就要保证所有其余的连接两个连通块的边的权值严格比这条边大$1$，只有这样，Kruskal算法才能得到唯一的一棵生成树。所以算法就是这样的：将树的$n-1$条边从小到大排序，然后枚举每条边$e$，求出它所连接的两个连通块各自的点数$x$和$y$，那么除了这条边之外，这两个连通块互相之间的边的长度应该恰好是$|e|+1$，那么答案要累加$(x+y-1)(|e|+1)$，直到所有边枚举完为止。代码如下：

#include <iostream>
#include <cstring>
#include <algorithm>
using namespace std;

const int N = 6010;
int n;
struct Edge {
    int a, b, w;
    bool operator<(const Edge & e) const {
        return w < e.w;
    }
} e[N];
int p[N], sz[N];

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

int main() {
    int T;
    cin >> T;
    while (T--) {
        cin >> n;
        for (int i = 0; i < n - 1; i++) {
            int a, b, w;
            cin >> a >> b >> w;
            e[i] = {a, b, w};
        }

        sort(e, e + n - 1);

        for (int i = 1; i <= n; i++) p[i] = i, sz[i] = 1;

        int res = 0;
        for (int i = 0; i < n - 1; i++) {
            int pa = find(p[e[i].a]), pb = find(p[e[i].b]);
            if (pa != pb) {
                res += (sz[pa] * sz[pb] - 1) * (e[i].w + 1);
                p[pa] = pb;
                sz[pb] += sz[pa];
            }
        }

        cout << res << endl;
    }

    return 0;
}

时间复杂度$O(n\log n)$，空间$O(n)$。