设
f
[
i
]
[
j
]
f[i][j]
f[i][j] 为离点
i
i
i 距离为
j
j
j 的点的点权和
因为这道题中点的状态跟祖先和儿子都有关系,所以我们要分两部转移
首先,对于点 x x x,我们只考虑在 x x x 的子树中距离为 j j j 的点,则转移方程为 f [ x ] [ j ] = ∑ f [ y ] [ j − 1 ] f[x][j]=\sum f[y][j-1] f[x][j]=∑f[y][j−1],其中 y y y 是 x x x 的儿子
然后,考虑用父亲来更新儿子:
f
[
y
]
[
j
]
=
f
[
y
]
[
j
]
+
f
[
x
]
[
j
−
1
]
−
f
[
y
]
[
j
−
2
]
f[y][j]=f[y][j]+f[x][j-1]-f[y][j-2]
f[y][j]=f[y][j]+f[x][j−1]−f[y][j−2]
由于是递归转移的,
f
[
x
]
[
j
−
1
]
f[x][j-1]
f[x][j−1] 在计算上面转移方程之前就已经算好了,所以当前
f
[
x
]
[
j
−
1
]
f[x][j-1]
f[x][j−1] 包含了所有与
x
x
x 的距离为
j
−
1
j-1
j−1 的点,而
f
[
y
]
[
j
−
2
]
f[y][j-2]
f[y][j−2] 和
f
[
y
]
[
j
]
f[y][j]
f[y][j] 只包含了
y
y
y 的子树中的点
仔细想想就会发现,
f
[
y
]
[
j
]
f[y][j]
f[y][j] 和
f
[
x
]
[
j
−
1
]
f[x][j-1]
f[x][j−1] 中有一个重合的部分(被算了两次),也就是
f
[
y
]
[
j
−
2
]
f[y][j-2]
f[y][j−2],所以我们要把它减掉
最后,点 i i i 对应的答案就是 ∑ j = 0 k f [ i ] [ j ] \sum\limits_{j=0}^{k} f[i][j] j=0∑kf[i][j]
关键代码
void work(int x,int fa) // 儿子更新父亲
{
for(int i=0;i<e[x].size();++i)
{
int y=e[x][i];
if(y==fa)continue;
work(y,x);
for(int j=1;j<=k;++j)
f[x][j]+=f[y][j-1];
}
}
void dfs(int x,int fa) // 父亲更新儿子
{
for(int i=0;i<e[x].size();++i)
{
int y=e[x][i];
if(y==fa)continue;
for(int j=k;j>=2;--j)
f[y][j]-=f[y][j-2];
for(int j=1;j<=k;++j)
f[y][j]+=f[x][j-1];
dfs(y,x);
}
}
以上是看题解后才get到的
开始我的思路跟题解有些不同,我的
f
[
i
]
[
j
]
f[i][j]
f[i][j] 为与点
i
i
i 距离不超过
j
j
j 的点的点权和
经过一番思考后,发现上面的转移方程仍然适用。我们只需要在 work()
函数后,通过一个简单的求和就可以了(本质上只是把最后的求和往前放了)。最后点
i
i
i 的答案就是
f
[
i
]
[
k
]
f[i][k]
f[i][k]
完整代码
#include<cstdio>
#include<iostream>
#include<vector>
using namespace std;
const int Maxn=100000+10,inf=0x3f3f3f3f;
int f[Maxn][23],a[Maxn];
int n,k;
vector <int> e[Maxn];
inline void add(int x,int y)
{
e[x].push_back(y);
e[y].push_back(x);
}
inline 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<<3)+(s<<1)+(ch^48),ch=getchar();
return s*w;
}
void work(int x,int fa)
{
for(int i=0;i<e[x].size();++i)
{
int y=e[x][i];
if(y==fa)continue;
work(y,x);
for(int j=1;j<=k;++j)
f[x][j]+=f[y][j-1];
}
}
void dfs(int x,int fa)
{
for(int i=0;i<e[x].size();++i)
{
int y=e[x][i];
if(y==fa)continue;
for(int j=k;j>=2;--j)
f[y][j]-=f[y][j-2];
for(int j=1;j<=k;++j)
f[y][j]+=f[x][j-1];
dfs(y,x);
}
}
int main()
{
// freopen("in.txt","r",stdin);
n=read(),k=read();
for(int i=1;i<n;++i)
{
int x=read(),y=read();
add(x,y);
}
for(int i=1;i<=n;++i)
f[i][0]=a[i]=read();
work(1,0); // 函数结束时,f[x][i] 只是包含子树中距离为i的点
for(int x=1;x<=n;++x)
for(int i=1;i<=k;++i)
f[x][i]+=f[x][i-1]; //求和后,f[x][i]就包含了子树中所有距离不超过i的点
dfs(1,0);
for(int i=1;i<=n;++i)
printf("%d\n",f[i][k]);
return 0;
}