[USACO12FEB]Nearby Cows G【换根DP】

题目链接

---

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

  求一个点的点权值很容易，以它为根，然后直接向下dfs，所有深度小于等于K的都是可以作为答案的，所以，处理子树上深度小于等于K的点权之和是有必要的，我们开个dp[maxN][maxK]来记录，每个点，向下kk步能积累的所有的点权之和。

  然后，再丢出一个问题，如果我们知道一个点的答案，我们是否可以知道它的一个邻接点的答案呢？当然是可以的，我们可以利用O(K)的办法来解决，用容斥的思想，不断的用父亲减去不可行，再加上一些可行却被删去的父亲的儿子的点权，具体看dp方程中的这段：

    ans[u] = ans[father] + dp[u][K];
    int ff = father, of = u;
    for(int kk = 1; kk <= K && ff; kk++)
    {
        ans[u] += dp[of][K - kk] - dp[ff][K - kk + 1];
        of = ff;
        ff = fa[ff];
    }
    if(ff) ans[u] -= a[ff];

  于是，我们换根的框架就搭好了。

#include <iostream>
#include <cstdio>
#include <cmath>
#include <string>
#include <cstring>
#include <algorithm>
#include <limits>
#include <vector>
#include <stack>
#include <queue>
#include <set>
#include <map>
#include <bitset>
//#include <unordered_map>
//#include <unordered_set>
#define lowbit(x) ( x&(-x) )
#define pi 3.141592653589793
#define e 2.718281828459045
#define INF 0x3f3f3f3f
#define eps 1e-8
#define HalF (l + r)>>1
#define lsn rt<<1
#define rsn rt<<1|1
#define Lson lsn, l, mid
#define Rson rsn, mid+1, r
#define QL Lson, ql, qr
#define QR Rson, ql, qr
#define myself rt, l, r
#define MP(a, b) make_pair(a, b)
using namespace std;
typedef unsigned long long ull;
typedef unsigned int uit;
typedef long long ll;
const int maxN = 1e5 + 7;
int N, K, head[maxN], cnt, a[maxN];
struct Eddge
{
    int nex, to;
    Eddge(int a=-1, int b=0):nex(a), to(b) {}
}edge[maxN << 1];
inline void addEddge(int u, int v)
{
    edge[cnt] = Eddge(head[u], v);
    head[u] = cnt++;
}
inline void _add(int u, int v) { addEddge(u, v); addEddge(v, u); }
int ans[maxN] = {0}, dp[maxN][21] = {0}, fa[maxN];
void pre_dfs(int u, int father)
{
    dp[u][0] = a[u]; fa[u] = father;
    for(int i=head[u], v; ~i; i=edge[i].nex)
    {
        v = edge[i].to;
        if(v == father) continue;
        pre_dfs(v, u);
        for(int kk=1; kk<=K; kk++) dp[u][kk] += dp[v][kk - 1];
    }
}
void dfs(int u, int father)
{
    ans[u] = ans[father] + dp[u][K];
    int ff = father, of = u;
    for(int kk = 1; kk <= K && ff; kk++)
    {
        ans[u] += dp[of][K - kk] - dp[ff][K - kk + 1];
        of = ff;
        ff = fa[ff];
    }
    if(ff) ans[u] -= a[ff];
    for(int i=head[u], v; ~i; i=edge[i].nex)
    {
        v = edge[i].to;
        if(v == father) continue;
        dfs(v, u);
    }
}
inline void init()
{
    cnt = 0;
    for(int i=1; i<=N; i++) head[i] = -1;
}
int main()
{
    scanf("%d%d", &N, &K);
    init();
    for(int i=1, u, v; i<N; i++)
    {
        scanf("%d%d", &u, &v);
        _add(u, v);
    }
    for(int i=1; i<=N; i++) scanf("%d", &a[i]);
    pre_dfs(1, 0);
    for(int kk=0; kk<=K; kk++) ans[1] += dp[1][kk];
    for(int i=head[1], v; ~i; i=edge[i].nex)
    {
        v = edge[i].to;
        dfs(v, 1);
    }
    for(int i=1; i<=N; i++) printf("%d\n", ans[i]);
    return 0;
}