[luoguP3047] [USACO12FEB]附近的牛Nearby Cows（DP）

传送门

dp[i][j][0] 表示点 i 在以 i 为根的子树中范围为 j 的解

dp[i][j][1] 表示点 i 在除去 以 i 为根的子树中范围为 j 的解

状态转移就很好写了

——代码

#include <cstdio>
#include <cstring>
#include <iostream>
#define N 100001

int n, k, cnt;
int f[N], dp[N][21][2], head[N], to[N << 1], next[N << 1], ans[N];

inline int read()
{
	int x = 0, f = 1;
	char ch = getchar();
	for(; !isdigit(ch); ch = getchar()) if(ch == '-') f = -1;
	for(; isdigit(ch); ch = getchar()) x = (x << 1) + (x << 3) + ch - '0';
	return x * f;
}

inline void add(int x, int y)
{
	to[cnt] = y;
	next[cnt] = head[x];
	head[x] = cnt++;
}

inline void dfs1(int u)
{
	int i, v;
	for(i = head[u]; i ^ -1; i = next[i])
	{
		v = to[i];
		if(v ^ f[u])
		{
			f[v] = u;
			dfs1(v);
			dp[v][0][1] = dp[u][0][0];
		}
	}
}

inline void dfs2(int u, int k)
{
	int i, v;
	dp[u][k][0] = dp[u][0][0];
	for(i = head[u]; i ^ -1; i = next[i])
	{
		v = to[i];
		if(v ^ f[u])
		{
			dfs2(v, k);
			dp[u][k][0] += dp[v][k - 1][0];

		}
	}
	for(i = head[u]; i ^ -1; i = next[i])
	{
		v = to[i];
		if(v ^ f[u]) dp[v][k][1] = dp[u][k - 1][1] + dp[u][k][0] - dp[v][k - 1][0];
	}
}

int main()
{
	int i, j, x, y;
	n = read();
	k = read();
	memset(head, -1, sizeof(head));
	for(i = 1; i < n; i++)
	{
		x = read();
		y = read();
		add(x, y);
		add(y, x);
	}
	for(i = 1; i <= n; i++) dp[i][0][0] = read();
	dfs1(1);
	for(i = 1; i <= k; i++) dfs2(1, i);
	for(i = 1; i <= n; i++) ans[i] = dp[i][k][0] + dp[i][k - 1][1];
	for(i = 1; i <= n; i++) printf("%d\n", ans[i]);
	return 0;
}