P3047 [USACO12FEB]附近的牛Nearby Cows

难得一见的换根法(肯定是我题目做少了吧)
首先考虑直接暴力dp, 设状态为 \(f[i][k]\), 表示第i个节点走k步的总数.
然后就能求出节点1的答案了.
可惜这样不够. 我们要求的是所有节点的答案. 而换一个节点作为根重新跑dp的话复杂度为 \(O(KN^2)\), 显然不现实.
考虑从已知答案的节点1推出其某个子节点(不妨设为2号节点)的答案.
明显, 假如以2号节点为根, 那么2节点的答案少了从1经过的节点数.
于是暴力加上去即可, 注意走到1算一步, 所以对于所有K, \(f[2][k] += f[1][k - 2]\).
还有一个小问题. 之前算1节点的时候, 2是1的子节点, 那么现在换根以后, 2实际上被算了2次, 所以还应该对于所有k, \(f[2][k] -= f[2][k - 1]\), 这里记得倒序.

#include <vector>
#include <cstdio>
#include <cstring>
#include <cassert>
#include <iostream>
#include <algorithm>
using namespace std;
const int MAXN = 1e5 + 10;
const int MAXK = 20 + 2;
inline int read(){
    char ch = getchar(); int x = 0;
    while(!isdigit(ch)) ch = getchar();
    while(isdigit(ch)) x = x * 10 + ch - '0', ch = getchar();
    return x;
}

int N, K;
int c[MAXN];
vector<int> g[MAXN];

int f[MAXN][MAXK];
void dp(int u, int fa) {
    f[u][0] = c[u];

    for(int i = 0; i < (int) g[u].size(); i++) {
        int &v = g[u][i];
        if(v == fa) continue;
        dp(v, u);
        for(int j = 1; j <= K; j++) f[u][j] += f[v][j - 1];
    }
}

int tot[MAXN];
void dfs(int u, int fa) {
    for(int i = 0; i <= K; i++) tot[u] += f[u][i];

    for(int i = 0; i < (int) g[u].size(); i++) {
        int &v = g[u][i];
        if(v == fa) continue;

        for(int j = K; j >= 2; j--) 
            f[v][j] += (f[u][j - 1] - f[v][j - 2]);
        f[v][1] += f[u][0];
        dfs(v, u);
    }
}

int main(){
    // freopen("p3047.in", "r", stdin);
    // freopen("p3047.out", "w", stdout);
    cin>>N>>K;
    for(int i = 1; i < N; i++) {
        int u = read(), v = read();
        g[u].push_back(v), g[v].push_back(u);
    }
    for(int i = 1; i <= N; i++) c[i] = read();
    dp(1, 1); dfs(1, 1);
    for(int i = 1; i <= N; i++) printf("%d\n", tot[i]);
    return 0;
}