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_JUDGE
    freopen("in.txt", "r", stdin);
    freopen("out.txt", "w", stdout);
    int StartTime = clock();
#endif
    scanf("%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_JUDGE
    printf("Run_Time = %d ms\n", clock() - StartTime);
#endif
    return 0;
}