51nod 1405 树的距离之和

这是一个树形dp的思想，两次dfs。

第一次dfs处理出，对于每个u，其所有子孙节点到其的距离之和。

第二次dfs处理出，对于每个u，所有点到其距离之和。

#include<iostream>
#include<string.h>
#include<string>
#include<algorithm>
#include<stdio.h>
#include<vector>
#include<set>
#include<time.h>
#define LL long long int
using namespace std;

const int maxn = 120000;
vector<int> G[maxn];
LL sum[maxn];
LL sz[maxn];
LL save[maxn];
int n;

void init()
{
    memset(G,0,sizeof(G));
    memset(sum,0,sizeof(sum));
    memset(sz,0,sizeof(sz));
    memset(save,0,sizeof(save));
}

LL dfs_1(int u,int fa)
{
    LL ans = 0;
    sz[u] = 1;
    for(int i = 0;i<G[u].size();i++)
    {
        int v = G[u][i];
        if(v == fa) continue;
        ans += dfs_1(v,u);
        ans += sz[v];
        sz[u] += sz[v];
    }
    sum[u] = ans;
    return ans;
}

void dfs_2(int u,int fa,LL up)
{
    save[u] = sum[u] + up;
    for(int i = 0;i<G[u].size();i++)
    {
        int v = G[u][i];
        if(v == fa) continue;
        dfs_2(v,u,up + n - sz[u] + sum[u] - sum[v] - sz[v] + sz[u] - sz[v]);
        /*
        next up = up + (n - sz[u]) + (sum[u] - sum[v] - sz[v]) + sz[u] - sz[v];
        up是原来的up,肯定要加
        对于u以上的所有节点，相当于所有的节点到他的距离+1,所以有n - sz[u]
        对于u的子树，sum[u] - sum[v] - sz[v]是指（除了v以外的，其他子树的所有节点到u的距离和)
        sz[u] - sz[v]是紧接着上一条的,因为u到v又推进了一条边，所以需要加一次长度
        */
    }
}


int main()
{
    while(cin >> n)
    {
        init();
        for(int i = 0;i<n-1;i++)
        {
            int u,v;
            scanf("%d %d",&u,&v);
            G[u].push_back(v);
            G[v].push_back(u);
        }
        dfs_1(1,-1);
        dfs_2(1,-1,0);
        for(int i = 1;i<=n;i++)
        {
            printf("%lld\n",save[i]);
        }
    }
    return 0;
}