51Nod 1405 - 树的距离之和(树DP)

在这里插入图片描述
【思路】
假设节点标号从0开始且0为树根,设 n u m [ u ] num[u] 表示记录包含 u u 在内 u u 的子节点个数, d i s t [ u ] dist[u] 记录 u u 的所有子节点到 u u 的距离之和, d p [ u ] dp[u] 记录最终答案,先计算所有的 n u m [ u ] num[u] d i s t [ u ] dist[u]
n u m [ u ] = 1 + n u m [ v ]       v u num[u]=1+\sum num[v] \ \ | \ \ v是u的子节点 d i s t [ u ] = d i s t [ v ] + n u m [ v ]       v u dist[u]=\sum dist[v]+num[v] \ \ | \ \ v是u的子节点 可以通过一次dfs计算出来,此时的 d i s t [ 0 ] = d p [ 0 ] dist[0]=dp[0] 就是答案,然后从根节点往下推 d p [ v ] = d p [ u ] ( n u m [ v ] ) + ( n n u m [ v ] ) = d p [ u ] + n 2 × n u m [ v ]       v u dp[v]=dp[u]-(num[v])+(n-num[v])=dp[u]+n-2×num[v] \ \ | \ \ v是u的子节点
减去 n u m [ v ] num[v] 是因为 v v 以及它的所有子节点到 v v 的距离比到 u u 的距离少1,同理减去 n n u m [ v ] n-num[v] 是因为剩余节点到 v v 的距离比到 u u 的距离多1

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int maxn=100005;

int n,rt,tot;
int head[maxn];
ll dist[maxn],num[maxn];
ll dp[maxn];
//num[u]记录包含u在内u的子节点个数,dist[u]记录u的所有子节点到u的距离之和

struct Edge{
    int to,next;
}e[maxn<<1];

void add(int u,int v){
    e[tot].to=v;
    e[tot].next=head[u];
    head[u]=tot++;
}

void dfs(int u,int fa){
    num[u]=1;
    dist[u]=0;
    for(int i=head[u];i!=-1;i=e[i].next){
        int v=e[i].to;
        if(v!=fa){
            dfs(v,u);
            num[u]+=num[v];
            dist[u]+=dist[v]+num[v];
        }
    }
}

void calc(int u,int fa){
    for(int i=head[u];i!=-1;i=e[i].next){
        int v=e[i].to;
        if(v!=fa){
            dp[v]=dp[u]+n-2*num[v];
            calc(v,u);
        }
    }
}

int main(){
    memset(head,-1,sizeof(head));
    scanf("%d",&n);
    for(int i=0;i<n-1;++i){
        int u,v;
        scanf("%d%d",&u,&v);
        --u;--v;
        add(u,v);
        add(v,u);
    }
    dfs(0,-1);
    dp[0]=dist[0];
    calc(0,-1);
    for(int i=0;i<n;++i) printf("%lld\n",dp[i]);
    return 0;
}

猜你喜欢

转载自blog.csdn.net/xiao_k666/article/details/88315585
今日推荐