[Luogu] P3047 [USACO12FEB]Nearby Cows G

\(Link\)

Description

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

Solution

注意到一个点只能被它的子树和上方的节点更新。

我们设\(dp1[x][i]\)表示\(x\)的子树内距离\(x\)恰好为\(i\)的节点的权值和，\(dp2[x][i]\)表示距离\(x\)恰好为\(i\)的节点的权值和。

\(dp1[x][i]\)是很好处理的，就是\(dp1[x][i]=\sum{dp1[y][i-1]}\)，然后\(dp1[x][0]=a[x]\)。

而\(dp2[x][i]\)不是很好处理。先让\(dp2[x][i]=dp1[x][i]\)，就只要再处理上方的贡献了。

那么我们用\(x\)来更新\(y\)，\(dp2[y][i]=dp2[y][i]+dp2[x][j-1]-dp2[y][j-2]\)（容斥，手动模拟可得）

最后的答案\(res[i]=\sum\limits_{j=0}^{k}dp2[i][j]\)。

Code

#include <bits/stdc++.h>

using namespace std;

int n, k, tot, g[20], res[100005], hd[100005], a[100005], dep[100005], nxt[200005], to[200005], dp1[100005][22], dp2[1000005][22];

int read()
{
	int x = 0, fl = 1; char ch = getchar();
	while (ch < '0' || ch > '9') { if (ch == '-') fl = -1; ch = getchar();}
	while (ch >= '0' && ch <= '9') {x = x * 10 + ch - '0'; ch = getchar();}
	return x * fl;
}

void add(int x, int y)
{
	tot ++ ;
	to[tot] = y;
	nxt[tot] = hd[x];
	hd[x] = tot;
	return;
}

void dfs1(int x, int fa)
{
	for (int i = hd[x]; i; i = nxt[i])
	{
		int y = to[i];
		if (y == fa) continue;
		dep[y] = dep[x] + 1;
		dfs1(y, x);
		for (int j = 1; j <= k; j ++ )
			dp1[x][j] += dp1[y][j - 1];
	}
	return;
}

void dfs2(int x, int fa)
{
	for (int i = hd[x]; i; i = nxt[i])
	{
		int y = to[i];
		if (y == fa) continue;
		for (int j = 2; j <= k; j ++ )
			g[j] = dp2[y][j - 2];
		for (int j = 1; j <= k; j ++ )
			dp2[y][j] = dp2[y][j] + dp2[x][j - 1] - g[j];
		dfs2(y, x);
	}
	return;
}

int main()
{
	n = read(); k = read();
	for (int i = 1; i <= n - 1; i ++ )
	{
		int x = read(), y = read();
		add(x, y); add(y, x);
	}
	for (int i = 1; i <= n; i ++ )
		a[i] = read(), dp1[i][0] = a[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 ++ )
		for (int j = 0; j <= k; j ++ )
			res[i] += dp2[i][j];
	for (int i = 1; i <= n; i ++ )
		printf("%d\n", res[i]);
	return 0;
}