牛客练习赛60 E.旗鼓相当的对手【树上启发式合并】

传送门

题解

这种统计所有子树的题肯定是可以用树上启发式合并做的。
预处理出每个点到根节点的距离，和重儿子。
然后开始做启发式合并，关键是如何统计每个节点的答案。
这里我参考了点分治的处理思路，在统计每个节点的答案之前，我一直了它一颗子树的所有信息，然后我就先遍历完它的一颗子树，将这颗子树的所有节点都取出来，然后统计一下答案，统计完之后再将它加入计数的数组，这样就可以避免计算同一颗子树上的点了。其实就是点分治的常规统计方法。
在即将向父节点回溯的时候，如果这个点是轻儿子，那么直接将计数数组中所有信息清除即可，当然不能memset，要遍历一下。如果这个点是重儿子，不要忘了它本身也应该被计数一下。

代码

#include <iostream>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <queue>
#define xx first
#define yy second
using namespace std;
typedef long long LL;
typedef pair<int,int> PII;
const int inf=0x3f3f3f3f;
const LL INF=0x3f3f3f3f3f3f3f3f;
const int N=3e5+10;
const int M=1e6+10;
int n,k,a[N],dep[N],siz[N],son[N];
LL ans[N],temp[N],cnt,suma[N],depcnt[N];
int head[N],to[N*2],nxt[N*2],tot;
void add(int u,int v) {to[++tot]=v;nxt[tot]=head[u];head[u]=tot;}

void predfs(int u,int fa){
	siz[u]=1;dep[u]=dep[fa]+1;
	for(int i=head[u];i;i=nxt[i]){
		int v=to[i];
		if(v==fa) continue;
		predfs(v,u);
		siz[u]+=siz[v];
		if(siz[v]>siz[son[u]]) son[u]=v;
	}
}

void getnode(int u,int fa){
	temp[++cnt]=u;
	for(int i=head[u];i;i=nxt[i]) if(to[i]!=fa) getnode(to[i],u);
}

void clear(int u,int fa){
	suma[dep[u]]=depcnt[dep[u]]=0;
	for(int i=head[u];i;i=nxt[i]) if(to[i]!=fa) clear(to[i],u);
}

void dfs(int u,int fa,bool keep){
	for(int i=head[u];i;i=nxt[i]){
		int v=to[i];
		if(v==fa||v==son[u]) continue;
		dfs(v,u,false);
	}
	if(son[u]) dfs(son[u],u,true);
	for(int ii=head[u];ii;ii=nxt[ii]){
		int v=to[ii];
		if(v==fa||v==son[u]) continue;
		cnt=0;
		getnode(v,u);
		for(int i=1;i<=cnt;i++)
			if(dep[temp[i]]-dep[u]<k){
				int depv=k+2*dep[u]-dep[temp[i]];
				ans[u]+=suma[depv]+depcnt[depv]*a[temp[i]];
			}
		for(int i=1;i<=cnt;i++) suma[dep[temp[i]]]+=a[temp[i]],depcnt[dep[temp[i]]]++;
	}
	suma[dep[u]]+=a[u],depcnt[dep[u]]++;
	if(!keep) clear(u,fa);
}

int main(){
	scanf("%d%d",&n,&k);
	for(int i=1;i<=n;i++) scanf("%d",&a[i]);
	for(int i=1,u,v;i<n;i++){
		scanf("%d%d",&u,&v);
		add(u,v);add(v,u);
	}
	predfs(1,0);
	dfs(1,0,true);
	for(int i=1;i<=n;i++) printf("%lld ",ans[i]);
	return 0;
} 

吐槽一下vector的DD速度，最好还是老老实实用链式向前星吧。