HDU6567树的重心换根DP

Posted 2021-09-07 hesorchen

题目

Cotree

用一条边将两颗树连接起来，使得点对之间的距离和最小。

求解

找出两棵树的重心，连接重心，换根DP求一下每条边的贡献即可。

代码

#include <bits/stdc++.h>
using namespace std;

const int N = 1e5 + 5;
vector<int> vec[N];
bool vis[N];
int nd1, nd2;
int N1, N2;
int siz[N];
void pre(int u, int fa)
{
    vis[u] = 1, ++N1;
    for (auto v : vec[u])
    {
        if (v == fa)
            continue;
        pre(v, u);
    }
}
int maxx;
int dfs1(int u, int fa)
{
    int mx = 0;
    for (auto v : vec[u])
    {
        if (v == fa)
            continue;
        int val = dfs1(v, u);
        siz[u] += val;
        mx = max(mx, val);
    }
    mx = max(mx, N1 - (siz[u] + 1));
    if (mx < maxx)
        nd1 = u, maxx = mx;
    return siz[u] + 1;
}
int dfs2(int u, int fa)
{
    int mx = 0;
    for (auto v : vec[u])
    {
        if (v == fa)
            continue;
        int val = dfs2(v, u);
        siz[u] += val;
        mx = max(mx, val);
    }
    mx = max(mx, N2 - (siz[u] + 1));
    if (mx < maxx)
        nd2 = u, maxx = mx;
    return siz[u] + 1;
}

long long ans;
int n;
long long dp(int u, int fa)
{
    for (auto v : vec[u])
    {
        if (v == fa)
            continue;
        int val = dp(v, u);
        ans += 1ll * val * (n - val);
        siz[u] += val;
    }
    return siz[u] + 1;
}

int main()
{
    ios::sync_with_stdio(false);
    cin.tie(nullptr), cout.tie(nullptr);
    cin >> n;
    for (int i = 1; i <= n - 2; i++)
    {
        int u, v;
        cin >> u >> v;
        vec[u].emplace_back(v);
        vec[v].emplace_back(u);
    }
    pre(1, -1);
    N2 = n - N1;
    maxx = 0x3f3f3f3f;
    memset(siz, 0, sizeof siz);
    dfs1(1, -1);
    maxx = 0x3f3f3f3f;
    memset(siz, 0, sizeof siz);
    for (int i = 1; i <= n; i++)
        if (!vis[i])
        {
            dfs2(i, -1);
            break;
        }
    vec[nd1].emplace_back(nd2);
    memset(siz, 0, sizeof siz);
    dp(1, -1);
    cout << ans << endl;
    return 0;
}