题目大意:给一个树,每次询问距离点(x)不超过(y)的所有点的点权和,(n le 10^5).
考虑树上( exttt{dp}),设(dp1_{i,j})为(i)子树下距离不超过(j)的所有点的点权和,易得:
[dp1_{i,j} = val_i + sum_{v in mathcal{S}} dp1_{v,j - 1}
]
考虑转换答案,设(dp2_{i,j})为(距离)i(不超过)j$的所有点的点权和,通过容斥得:
[dp2_{i,j} = dp1_{i,j} + dp2_{fa,j - 1} - [j ge 2]dp2_{i,j - 2}
]
其中(fa)为(i)的父亲节点。
#include <bits/stdc++.h>
using namespace std;
const int N = 1e5 + 5;
int n, val[N], k;
vector<int> g[N];
int dp1[N][22], dp2[N][22];
void dfs1(int x, int fa)
{
for (int i = 0; i <= k; i++)
dp1[x][i] = val[x];
for (int i = 0; i < g[x].size(); i++)
{
int v = g[x][i];
if (v != fa)
{
dfs1(v, x);
for (int j = 1; j <= k; j++)
{
dp1[x][j] += dp1[v][j - 1];
}
}
}
return;
}
void dfs2(int x, int fa)
{
for (int i = 0; i < g[x].size(); i++)
{
int v = g[x][i];
if (v != fa)
{
dp2[v][1] += dp1[x][0];
for (int j = 2; j <= k; j++)
{
dp2[v][j] += dp2[x][j - 1] - dp1[v][j - 2];
}
dfs2(v, x);
}
}
return;
}
int main(void)
{
scanf("%d%d", &n, &k);
for (int i = 1; i < n; i++)
{
int u, v;
scanf("%d%d", &u, &v);
g[u].push_back(v);
g[v].push_back(u);
}
for (int i = 1; i <= n; i++)
scanf("%d", &val[i]);
dfs1(1, 0);
for (int i = 1; i <= n; i++)
{
for (int j = 0; j <= k; j++)
dp2[i][j] = dp1[i][j];
}
dfs2(1, 0);
for (int i = 1; i <= n; i++)
{
printf("%d
", dp2[i][k]);
}
return 0;
}