BZOJ 3011 Running Away From the Barn

题目大意

给你一棵树，树上的边有边权，求每个节点子树中距离不超过k的节点个数。

这题考虑DFS维护子树信息。

因为距离是固定的，所以如果一个点和它一个儿子的距离超过k那么它祖先也一定符合这个要求。所以可以使用可并堆进行维护。

每次进行合并之后去除所有距离大于k的点，之后的就是这个点的答案。因为每个点只进堆一次出堆一次所以时间复杂度是$O(N \log N)$的。

#include <cstdio>
#include <algorithm>
using namespace std;
#define N 200010

int n,fa[N],cnt[N],rt[N],ls[N],rs[N],h[N];
int to[N],nex[N],head[N];
long long l,val[N],depth[N];

int merge(int tx,int ty)
{
	if(!tx) return ty;
	if(!ty) return tx;
	if(depth[tx]<depth[ty]) swap(tx,ty);
	rs[tx]=merge(rs[tx],ty);
	if(h[ls[tx]]<h[rs[tx]]) swap(ls[tx],rs[tx]);
	h[tx]=rs[tx]?h[rs[tx]]+1:0;
	return tx;
}

void DFS(int pos)
{
	for(int i=head[pos];i;i=nex[i])
	{
		DFS(to[i]);
		rt[pos]=merge(rt[to[i]],rt[pos]);
		cnt[pos]+=cnt[to[i]];
	}
	while(depth[rt[pos]]>depth[pos]+l) rt[pos]=merge(ls[rt[pos]],rs[rt[pos]]),cnt[pos]--;
}

void init(int pos)
{
	for(int i=head[pos];i;i=nex[i])
	{
		depth[to[i]]=depth[pos]+val[i];
		init(to[i]);
	}
}

int main()
{
	scanf("%d%lld",&n,&l);
	for(int i=2;i<=n;++i)
	{
		scanf("%d",fa+i);
		to[i]=i;
		nex[i]=head[fa[i]];
		head[fa[i]]=i;
		scanf("%lld",val+i);
	}
	init(1);
	for(int i=1;i<=n;++i) cnt[i]=1,rt[i]=i;
	DFS(1);
	for(int i=1;i<=n;++i) printf("%d\n",cnt[i]);
	return 0;
}