\(\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\n", ans);
}
return 0;
}