2019CCPC-江西省赛 Cotree（树的重心 并查集）

传送门

题意

给两棵树，在两棵树中连一条边，
使
${\sum }_{i=1}^n{\sum }_{j=i+1}^{n}dis(i,j)$最小。

分析

需要找到两棵树的重心，将两棵树的重心连起来，再统计每条边对答案的贡献。
对于输入的处理用并查集即可。

代码

#include <bits/stdc++.h>using namespace std;
//-----pre_def----
const double PI = acos(-1.0);
const int INF = 0x3f3f3f3f;
typedef long long LL;
typedef unsigned long long ULL;
typedef pair<int, int> PII;
typedef pair<double, double> PDD;
#define fir(i, a, b) for (int i = (a); i <= (b); i++)
#define rif(i, a, b) for (int i = (a); i >= (b); i--)
#define endl '\n'
#define init_h memset(h, -1, sizeof h), idx = 0;
#define lowbit(x) x &(-x)//---------------
const int N = 1e5 + 10;
int n;
int h[N], e[N << 1], ne[N << 1], idx;
int p[N];
int root1, root2, sz[N], sum[N], w[N];
void init()
{
    init_h;fir(i, 1, n) p[i] = i, sum[i] = 1;w[0] = INF;
}void add(int a, int b)
{
    e[idx] = b;ne[idx] = h[a];h[a] = idx++;
}
int find(int x)
{
    if (p[x] != x)p[x] = find(p[x]);return p[x];
}void dfs1(int u, int fa)
{
    sz[u] = 1;w[u] = 0;for (int i = h[u]; ~i; i = ne[i]){
    int t = e[i];if (t == fa)continue;dfs1(t, u);sz[u] += sz[t];w[u] = max(w[u], sz[t]);}w[u] = max(w[u], sum[find(u)] - sz[u]);if (w[u] <= w[root1]) //重心的性质{
    root1 = u; //可能有两个重心}
}
void dfs2(int u, int fa)
{
    sz[u] = 1;w[u] = 0;for (int i = h[u]; ~i; i = ne[i]){
    int t = e[i];if (t == fa)continue;dfs2(t, u);sz[u] += sz[t];w[u] = max(w[u], sz[t]);}w[u] = max(w[u], sum[find(u)] - sz[u]);if (w[u] <= w[root2]) //重心的性质{
    root2 = u; //可能有两个重心}
}
LL ans = 0;
void get_ans(int u, int fa)
{
    sz[u] = 1;for (int i = h[u]; ~i; i = ne[i]){
    int t = e[i];if (t == fa)continue;get_ans(t, u);sz[u] += sz[t];ans += ((1ll * sz[t]) * (LL)(n - sz[t]));}
}
int main()
{
    
#ifndef ONLINE_JUDGEfreopen("in.txt", "r", stdin);freopen("out.txt", "w", stdout);int StartTime = clock();
#endifscanf("%d", &n);init();fir(i, 1, n - 2){
    int a, b;scanf("%d%d", &a, &b);add(a, b);add(b, a);int fa = find(a);int fb = find(b);if (fa != fb){
    p[fa] = fb;sum[fb]+=sum[fa];}}dfs1(1, -1);for (int i = 1; i <= n; i++){
    if (sz[i] == 0){
    dfs2(i, -1);break;}}add(root1, root2);add(root2, root1);get_ans(1, -1);printf("%lld\n", ans);
#ifndef ONLINE_JUDGEprintf("Run_Time = %d ms\n", clock() - StartTime);
#endifreturn 0;
}