51NOD-1405 树的距离之和（树形DP）

1405 树的距离之和

$$基准时间限制：1 秒\ \ 空间限制：131072 KB\ \ 分值: 40\ \ 难度：4级算法题$$

Description

$给定一棵无根树，假设它有n个节点，节点编号从1到n，$求任意两点之间的距离（最短路径）之和。

Input

$第一行包含一个正整数n (n <= 100000)，表示节点个数。 后面(n - 1)行，$每行两个整数表示树的边。

Output

$每行一个整数，第i(i = 1,2,…n)行表示所有节点到第i个点的距离之和。$

Sample Input

1 2
3 2
4 2

Sample Output

5
3
5
5

思路

首先可以想到在$O(n)$内用树形DP求出某一点到其他所有点的距离和。
设$dp[u]表示以结点u到其子树所有结点的距离和，cnt[u]表示包括u$在内的子树中所有结点个数，$v为u的子结点$，则有：

$$\begin{eqnarray*} cnt[u]&=&(\sum cnt[v])+1\\ dp[u]&=&\sum (dp[v]+cnt[v]) \end{eqnarray*}$$

$所以可以从叶子结点往上进行状态转移，直至算出根节点的dp值 又注意到若dp[u]表示u$到其他所有结点的距离和，$v为u的子结点$，则有：

$$\begin{eqnarray*} dp[v]&=&dp[u]-cnt[v]+(n-cnt[v])\\ &=&dp[u]+n-2*cnt[v]\\ \end{eqnarray*}$$
$v子树中的所有结点的到v的距离比其到u$的距离的各少 $1，共cnt[v]$个结点，则需要减去 $cnt[v]$
$不在v子树中的所有结点到v的距离比其到u的距离各多1，共n-cnt[v]$个结点，则需要加上 $n-cnt[v]$
$则可以在O(n)内算出所有的dp值$
$总的时间复杂度为O(n)$

代码

#pragma comment(linker, "/STACK:10240000,10240000")
#include <cstdio>
#include <cstring>
#include <vector>
#include <algorithm>

using namespace std;

int n,s,e;
int cnt[100005];//cnt[i]表示包括i在内的子树中所有结点个数
long long dp[100005];
vector<int> g[100005];

void dp_up(int u,int p) {//根据子结点的dp求父结点的dp，此时dp[i]表示i到其子树上所有结点的距离和
    cnt[u]=1;//先把自己算入
    int v;
    for(int i=g[u].size()-1;i>=0;--i) {
        v=g[u][i];
        if(v!=p) {
            dp_up(v,u);
            cnt[u]+=cnt[v];
            dp[u]+=dp[v]+cnt[v];//u到v子树上某一结点的距离 比 v到其子树上该结点的距离多1，v子树上共有cnt[v]个结点，则距离共多cnt[v]
        }
    }
}

void dp_down(int u,int p) {//根据父结点的dp求子结点的dp，此时dp[i]表示i到其他所有结点的距离和
    int v;
    for(int i=g[u].size()-1;i>=0;--i) {
        v=g[u][i];
        if(v!=p) {
            dp[v]=dp[u]+n-2*cnt[v];
            dp_down(v,u);
        }
    }
}

int main() {
    while(1==scanf("%d",&n)) {
        for(int i=1;i<n;++i) {
            scanf("%d%d",&s,&e);
            g[s].push_back(e);
            g[e].push_back(s);
            dp[i]=0;
        }
        dp[n]=0;
        dp_up(1,0);//求出其余点到结点1的距离和
        dp_down(1,0);//求出任意结点到其他所有点的距离和
        for(int i=1;i<=n;++i) {
            printf("%I64d\n",dp[i]);
        }
    }
    return 0;
}