P3478 [POI2008] STA-Station(换根DP 二次扫描)

P3478 [POI2008] STA-Station

由题意可得，要求以某点为根的深度之和最大，先假设1为根，f[i]表示以1为根时i的子树中的节点到i的贡献和，即f[i]=所有儿子节点的贡献值+儿子节点的个数，sz[i]表示的是以i为根的子树的大小，那么就很容易计算出f[1]的深度之和。这是以1为根的时候，那现在考虑如果把根从1移到1的儿子节点会发生什么，假设根从1移动到k，其中k是1的儿子，容易知道的是k的儿子对k的贡献大小是不改变的，即f[k]不会发生变化，发生变化的是除了k的子树外，其他所有节点对k的贡献，令v[i]表示当根为i时，除了i的子树外，其他所有节点对i的贡献之和，可知的是一开始我们就以1为根，然后把根往下放，所以1节点是没有父亲节点的，所以v[1]=0。考虑到这，我们再考虑v[k]怎么求，其中f[k]保持不变，即v[k]中不需要计算来自k节点子树的贡献值，只需要考虑除子树外其他所有节点的贡献值，即节点1与它的其他所有子树，可知除k节点外，其他所有子树对1的贡献为f[1]-f[k]-sz[k]，那么当1及其他所有子树对根产生贡献时，需要再加上一个节点个数，因此v[k]=f[1]-f[k]-sz[k]+n-sz[k]，这是以1为根，讲根下放到1的儿子节点的情况，再一般的，当以i为根，将根下放到i的儿子节点j的时候，如何将v[i]转移到v[j]呢？求解也很相似，不难发现的是f[j]依然保持不变，当以i为根的时候所有节点对i的贡献为f[i]+v[i]，所以求解除了j的子树外其他所有节点对j的贡献就是v[j]=f[i]+v[i]-f[j]-sz[j]+n-sz[j]，这样一定会计算到所有节点为根时产生的贡献吗？答案是肯定的，因为一开始我们计算贡献的时候是以1为根的，当我们再把根下放、计算贡献的时候，儿子节点的父亲节点是已经计算完的，那么转移的时候一定是从一个计算完贡献的根转移到一个还未计算贡献的儿子节点上，这样保证了一定能计算完以所有的节点为根时产生贡献的情况，时间复杂度O(n)

AC代码：

#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
int n;
LL f[1000010], sz[1000010], v[1000010];
vector<int> G[1000010];
void dfs1(int u, int fa) {
    sz[u] = 1;
    for (auto v : G[u]) {
        if (v != fa) {
            dfs1(v, u);
            sz[u] += sz[v];
            f[u] += f[v];
        }
    }
    f[u] += sz[u] - 1;
}
void dfs2(int u, int fa) {
    if (u == 1) {
        v[u] = 0;
    } else {
        v[u] = v[fa] + f[fa] - f[u] + n - 2 * sz[u];
    }
    for (auto v : G[u]) {
        if (v != fa) {
            dfs2(v, u);
        }
    }
}
int main() {
    scanf("%d", &n);
    for (int i = 1; i < n; i++) {
        int u, v;
        scanf("%d%d", &u, &v);
        G[u].push_back(v);
        G[v].push_back(u);
    }
    dfs1(1, 0);
    dfs2(1, 0);
    LL ans = v[1] + f[1];
    int id = 1;
    for (int i = 2; i <= n; i++) {
        if (v[i] + f[i] > ans) {
            ans = v[i] + f[i];
            id = i;
        }
    }
    printf("%d\n", id);
    return 0;
}