**题意**
对一颗$n$个节点的点权树，对于每个距离为k的点对$(u,v)$，如果
$LCA(u,v)!=u&&LCA(u,v)!=v$
则节点$LCA(u,v)$会获得$a_{u}+a_{v}$的贡献。
需要求出每个节点的贡献。
设$d[i]$为i节点的深度。
可以想到对于子树$x$中某个节点$y$,则$y$可以和$x$的其他子树中深度为$2*d[x]+k-d[y]$的点对$x$产生贡献。
这种找每个点贡献的题目自然会联想到dsu on tree。
因为所求与深度有关，所以也可以用长链剖分。
借这道题重温一下这两种解法
**dsu one tree**
dsu的大致流程如下：
PS:先重链剖分，预处理出每个节点的重儿子。
主流程：
1.对于节点$i$，遍历递归解决所有的轻儿子子树。
2.递归计算重儿子子树。
3.遍历每个轻儿子，计算轻儿子子树对与重儿子子树的贡献。
4.将轻儿子子树的状态加加入到重儿子子树。
5.当前节点$i$如果是轻儿子，则计算完成后从重儿子子树中删除轻儿子子树的状态。
大致代码如下
1 void dsu(int x,int f){ 2 for(int i=head[x];i!=-1;i=edge[i].next){ 3 int y = edge[i].e; 4 if(y==fat[x]||y==son[x])continue; 5 dsu(y,0); 6 }//暴力计算轻儿子 7 if(son[x])dsu(son[x],1);//计算重儿子 8 for(int i=head[x];i!=-1;i=edge[i].next){ 9 int y = edge[i].e; 10 if(y==fat[x]||y==son[x])continue; 11 ans[x]+=cal(y);//计算轻儿子对点x的贡献 12 add(y);//添加轻儿子状态 13 } 14 XXXXXXX;//添加重儿子状态 15 if(!f)del(x);//删除轻儿子状态 16 }
粗略一看全是暴力，为什么可行呢？
那么来分析下时间复杂度：
我们考虑一个点会被访问多少次
1.每个节点被访问，要么是它的祖先节点暴力统计轻儿子/消除影响的时候，要么是它自己统计答案的时候。
2.对于每个节点到根节点的路径上的轻边数量不会超过$logn$。
所以复杂度为$O(n*logn)$
#include<bits/stdc++.h> using namespace std; typedef long long ll; const int maxn = 2e5+10; const ll inf = 1e18; struct node { int s,e,next; } edge[maxn*2]; int head[maxn],len; void init() { memset(head,-1,sizeof(head)); len=0; } void add(int s,int e) { edge[len]= {s,e,head[s]}; head[s]=len++; } int son[maxn],siz[maxn],d[maxn],fat[maxn]; void dfs(int x,int fa) { son[x]=0,siz[x]=1; for(int i=head[x]; i!=-1; i=edge[i].next) { int y =edge[i].e; if(y==fa)continue; fat[y]=x; d[y]=d[x]+1; dfs(y,x); siz[x]+=siz[y]; if(!son[x]||siz[son[x]]<siz[y]) son[x]=y; } } int n,k; ll ans[maxn],cnt[maxn],num[maxn],a[maxn]; ll cal(int x,int lcad){ ll ans =0; if(k+2*lcad-d[x]>=0) ans+=cnt[2*lcad+k-d[x]]+a[x]*num[2*lcad+k-d[x]]; for(int i=head[x];i!=-1;i=edge[i].next){ int y = edge[i].e; if(y==fat[x])continue; ans+=cal(y,lcad); } return ans; } void add(int x){ cnt[d[x]]+=a[x],num[d[x]]++; for(int i=head[x];i!=-1;i=edge[i].next){ int y = edge[i].e; if(y==fat[x])continue; add(y); } } void del(int x){ cnt[d[x]]-=a[x],num[d[x]]--; for(int i=head[x];i!=-1;i=edge[i].next){ int y = edge[i].e; if(y==fat[x])continue; del(y); } } void dsu(int x,int f){ for(int i=head[x];i!=-1;i=edge[i].next){ int y = edge[i].e; if(y==fat[x]||y==son[x])continue; dsu(y,0); } if(son[x])dsu(son[x],1); for(int i=head[x];i!=-1;i=edge[i].next){ int y = edge[i].e; if(y==fat[x]||y==son[x])continue; ans[x]+=cal(y,d[x]); add(y); } cnt[d[x]]+=a[x],num[d[x]]++; if(!f)del(x); } int main() { init(); scanf("%d%d",&n,&k); for(int i=1; i<=n; i++) scanf("%d",&a[i]); for(int i=1,x,y; i<n; i++) { scanf("%d%d",&x,&y); add(x,y); add(y,x); } dfs(1,0); dsu(1,1); for(int i=1; i<=n; i++) printf("%lld ",ans[i]); }
**长链剖分**
长链剖分与重链剖分的不同之处为预处理出来的是深度最大的重儿子
大致流程如下：
PS:先长链剖分，预处理出每个节点的重儿子。
主流程：
1.对于节点$i$，递归计算重儿子子树。
2.遍历每个轻儿子，计算轻儿子子树对与重儿子子树的贡献。
4.将轻儿子子树的信息暴力合并到重儿子子树。
大致代码如下
1 void dfs2(int x,int top) { 2 if(son[x]) 3 dfs2(son[x],top); 4 for(int i=head[x]; i!=-1; i=edge[i].next) { 5 int y = edge[i].e; 6 if(y==fat[x]||y==son[x])continue; 7 dfs2(y,y); 8 XXXXXX//计算x节点的答案 9 XXXXXX//轻儿子子树状态合并到x中 10 } 11 }
同样分析一下复杂度：
对于节点$i$，可以直接从重儿子继承状态，暴力合并轻儿子
因为每个节点仅属于一条长链，且一条长链只会在链顶位置作为轻儿子暴力合并一次。
时间复杂度为$O(n)$
1 #include<bits/stdc++.h> 2 using namespace std; 3 typedef long long ll; 4 const int maxn = 2e5+10; 5 const ll inf = 1e18; 6 struct node { 7 int s,e,next; 8 } edge[maxn*2]; 9 int head[maxn],len; 10 void init() { 11 memset(head,-1,sizeof(head)); 12 len=0; 13 } 14 void add(int s,int e) { 15 edge[len]= {s,e,head[s]}; 16 head[s]=len++; 17 } 18 int son[maxn],fd[maxn],d[maxn],fat[maxn]; 19 void dfs1(int x,int fa) { 20 for(int i=head[x]; i!=-1; i=edge[i].next) { 21 int y =edge[i].e; 22 if(y==fa)continue; 23 fat[y]=x; 24 dfs1(y,x); 25 if(fd[y]>fd[son[x]]) 26 son[x]=y; 27 } 28 fd[x]=fd[son[x]]+1; 29 } 30 int n,k,a[maxn],num[maxn],nums; 31 ll ans[maxn]; 32 vector<ll>dp[maxn][2]; 33 void dfs2(int x,int top) { 34 if(x==top) { 35 num[x]=++nums; 36 dp[num[x]][0].resize(fd[x],0); 37 dp[num[x]][1].resize(fd[x],0); 38 } 39 if(son[x]) { 40 num[son[x]]=num[x]; 41 d[son[x]]=d[x]+1; 42 dfs2(son[x],top); 43 } 44 dp[num[x]][0][d[x]]++; 45 dp[num[x]][1][d[x]]+=a[x]; 46 for(int i=head[x]; i!=-1; i=edge[i].next) { 47 int y = edge[i].e; 48 if(y==fat[x]||y==son[x])continue; 49 dfs2(y,y); 50 for(int j = 0; j < fd[y]&&j<k-1; j++) 51 if(k - j - 1 + d[x] < fd[top]) { 52 ans[x] += dp[num[x]][0][k - j - 1 + d[x]] * dp[num[y]][1][j]; 53 ans[x] += dp[num[x]][1][k - j - 1 + d[x]] * dp[num[y]][0][j]; 54 } 55 for(int j=0; j<fd[y]; j++) { 56 dp[num[x]][0][j+1+d[x]]+=dp[num[y]][0][j]; 57 dp[num[x]][1][j+1+d[x]]+=dp[num[y]][1][j]; 58 } 59 } 60 } 61 int main() { 62 init(); 63 scanf("%d%d",&n,&k); 64 for(int i=1; i<=n; i++) 65 scanf("%d",&a[i]); 66 for(int i=1,x,y; i<n; i++) { 67 scanf("%d%d",&x,&y); 68 add(x,y); 69 add(y,x); 70 } 71 dfs1(1,0); 72 dfs2(1,1); 73 for(int i=1; i<=n; i++) 74 printf("%lld ",ans[i]); 75 }