P3047 [USACO12FEB]附近的牛Nearby Cows

(color{#0066ff}{ 题目描述 })

题目简述：给出一棵n个点的树，每个点上有(C_i)头牛，问每个点k步范围内各有多少头牛。

(color{#0066ff}{输入格式})

第一行：n和k;

后面n-1行：i和j（两块田野）;

之后n行：1..n每一块的C(i);

(color{#0066ff}{输出格式})

n行：每行M(i);//i：1..2

(color{#0066ff}{输入样例})

6 2 
5 1 
3 6 
2 4 
2 1 
3 2 
1 
2 
3 
4 
5 
6 

(color{#0066ff}{输出样例})

15 
21 
16 
10 
8 
11 

(color{#0066ff}{数据范围与提示})

1 <= n <= 100000,1<=k<=20

(color{#0066ff}{ 题解 })

设(f[i][j])表示距i距离为j的点权和(f[i][0]=val[i])

首先，dfs一遍可以求出子树内的节点对它的贡献

之后考虑父亲对自己的贡献，显然父亲的j-1会对自己的j有贡献

不过父亲的j-1会包括自己子树内的，有重复，这部分就是自己的j-2，容斥减一下就行了

#include<bits/stdc++.h>
#define LL long long
LL in() {
	char ch; LL x = 0, f = 1;
	while(!isdigit(ch = getchar()))(ch == '-') && (f = -f);
	for(x = ch ^ 48; isdigit(ch = getchar()); x = (x << 1) + (x << 3) + (ch ^ 48));
	return x * f;
}
struct node {
	int to;
	node *nxt;
	node(int to = 0, node *nxt = NULL): to(to), nxt(nxt) {}
	void *operator new (size_t) {
		static node *S = NULL, *T = NULL;
		return (S == T) && (T = (S = new node[1024]) + 1024), S++;
	}
};
const int maxn = 1e5 + 100;
LL f[maxn][25], val[maxn];
node *head[maxn];
int n, k;
void add(int from, int to) {
	head[from] = new node(to, head[from]);
}
void dfs(int x, int fa) {
	f[x][0] = val[x];
	for(node *i = head[x]; i; i = i->nxt) {
		if(i->to == fa) continue;
		dfs(i->to, x);
		for(int j = 1; j <= k; j++) f[x][j] += f[i->to][j - 1];
	}
}
void dfs2(int x, int fa) {
	for(node *i = head[x]; i; i = i->nxt) {
		if(i->to == fa) continue;
		for(int j = k; j >= 2; j--) f[i->to][j] -= f[i->to][j - 2];
		for(int j = 1; j <= k; j++) f[i->to][j] += f[x][j - 1];
		dfs2(i->to, x);
	}
}
int main() {
	n = in(), k = in();
	int x, y;
	for(int i = 1; i < n; i++) {
		x = in(), y = in();
		add(x, y), add(y, x);
	}
	for(int i = 1; i <= n; i++) val[i] = in();
	dfs(1, 0);
	dfs2(1, 0);
	for(int i = 1; i <= n; i++) {
		LL ans = 0;
		for(int j = 0; j <= k; j++) ans += f[i][j];
		printf("%lld
", ans);
	}
	return 0;
}