luogu p3047题解

题目大意

给你一颗$n$个点的树，点带权，对于每个节点求出距离它不超过$k$的所有节点的权值和。

题目思路

考虑到$k$比较小，所以我们先预处理$f{1}_{u,i}$，表示$u$子树中距离它不超过$i$的点的权值和。

然后进行换根$dp$。

从$u$到$v$，距离加$1$。

设$f{2}_{u,i}$表示$u$子树外距离它不超过$i$的点的权值和。

所以$f{2}_{v,i}=f{2}_{u,i-1}+f{1}_{u,j-1}-f{1}_{v,j-2}$。

也就是$u$的贡献加上他兄弟的贡献。

具体实现参考代码。

#include
using namespace std;
const int N=1e5+10;
int n,k,a[N],f1[N][20+10],f2[N][20+10];
int head[N<<1],nxt[N<<1],to[N<<1],tot=0;
void add(int u,int v)
{
    to[++tot]=v;
    nxt[tot]=head[u];
    head[u]=tot;
}
int read()
{
    int s=0,w=1;
    char ch=getchar();
    while(ch<'0'||ch>'9')
    {
        if(ch=='-')
            w=-1;
        ch=getchar();
    }
    while(ch>='0'&&ch<='9')
        s=s*10+(ch-'0'),ch=getchar();
    return s*w;
}
void dfs1(int u,int fa)
{
    for(int i=0;i<=k;++i)
        f1[u][i]=a[u];
    for(int i=head[u];i;i=nxt[i])
    {
        int v=to[i];
        if(v==fa)
            continue;
        dfs1(v,u);
        for(int j=1;j<=k;++j)
            f1[u][j]+=f1[v][j-1];
    }
}
void dfs2(int u,int fa)
{
    for(int i=head[u];i;i=nxt[i])
    {
        int v=to[i];
        if(v==fa)
            continue;
        f2[v][1]=f2[u][0]+f1[u][0];
        for(int j=2;j<=k;++j)
            f2[v][j]=f2[u][j-1]+f1[u][j-1]-f1[v][j-2];
        dfs2(v,u);            
    }
}
int main()
{
    n=read(),k=read();
    for(int i=1;i<=n-1;++i)
    {
        int u=read(),v=read();
        add(u,v);
        add(v,u);
    }
    for(int i=1;i<=n;++i)
        a[i]=read();
    dfs1(1,1);
    dfs2(1,1);
    for(int i=1;i<=n;++i)
        printf("%d\n",f1[i][k]+f2[i][k]);
    return 0;
}