P3047 [USACO12FEB]Nearby Cows G

题目大意：给一个树，每次询问距离点\(x\)不超过\(y\)的所有点的点权和，\(n \le 10^5\).

考虑树上\(\texttt{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\n", dp2[i][k]);
    }
    return 0;
}