hhc 1742 树的距离之和 树形dp

题目链接戳这里

时间限制:10000ms
单点时限:1000ms
内存限制:256MB
描述
给定一棵包含N个节点的树，节点编号1~N。  

请你对于每一个节点，计算它到所有其他节点的距离之和。(每条边的长度为1)

输入
第一行包含一个整数N。  

以下N-1行，每行包含两个整数a和b，代表ab之间有一条边。  

对于30%的数据，1 ≤ N ≤ 1000  

对于100%的数据，1 ≤ N ≤ 100000

输出
输出N行，依次对于1~N号节点输出答案。

样例输入
4  
1 2  
2 3  
2 4
样例输出
5  
3  
5  
5

我们任取节点1为树根，分两次dfs，第一次求节点i形成的子树的距离之和。第二次dfs求节点i到其它所有节点的距离和。
dfs1：根据子节点v的dp求父节点u的dp，此时dp[i]表示i到其子树上所有节点距离和。cnt是子节点的结点数量之和(包括u)。因为子树里的所有节点到v的距离，都会比到u的距离少1，所以得补上，这个差值是1*cnt[v]。所以
dp[u]+=dp[v]+cnt[v]

dfs2：根据父节点u的dp求子节点的dp，此刻dp[i]是i到其它所有结点的距离和。我们可以参考下面这幅图。

其中方框代表树，dp[u]是u到其子树所有节点的距离之和。白树（白色方框代表的树）内节点到v的距离相当于白树到u的距离-1，数量有cnt[v]个，所以dp[u]-cnt[v]，v到橙树的距离，相当于u到橙树距离+1，数量有N-cnt[v]个，所以dp[u]还得加上1*(N-cnt[v])
综上所述,dp[v]=dp[u]-cnt[v]+(N-cnt[v])
得代码:

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

#define pii pair<int,int>
#define ll long long
#define mst(a,b) memset(a,b,sizeof(a))
#define rep(i,a,b) for(ll i=(a);i<(b);++i)
#define pb push_back
const int inf = 0x3f3f3f3f, maxN = 1e5 + 5;
int N, M, T;
vector<int> G[maxN];
ll cnt[maxN], dp[maxN];

ll dfs1(int u, int f) {
    cnt[u] = 1;
    rep(i, 0, G[u].size()) {
        int v = G[u][i];
        if (v == f)
            continue;
        dfs1(v, u);
        cnt[u] += cnt[v];
        dp[u] += dp[v] + cnt[v];
    }
}

ll dfs2(int u, int f) {
    rep(i, 0, G[u].size()) {
        int v = G[u][i];
        if (v == f)
            continue;
        dp[v] = dp[u] - cnt[v] + (N - cnt[v]);
        dfs2(v, u);
    }
}

int main() {
    scanf("%d", &N);
    int u, v;
    rep(i, 0, N - 1) {
        scanf("%d%d", &u, &v);
        G[u].pb(v);
        G[v].pb(u);
    }
    dfs1(1, 0);
    dfs2(1, 0);
    rep(i, 1, N + 1)
        printf("%lld\n", dp[i]);
    return 0;
}