BZOJ 3653: 谈笑风生(主席树)

传送门

解题思路

　　首先对于一个$a$来说，要求$b$和$c$，那么$a,b,c$一定在一条链上。把$b$分类讨论，如果$b$是$a$的祖宗，这个方案数就很好统计了，就是$c$在$a$的子树里随便选，产生的贡献为$(siz_a-1)*(min(k,dep_a))$。如果$b$是$a$的儿子，那么就考虑$b$作为每个点产生的贡献，发现为$siz_x-1$，那么其实要求的就是$[dfn_a+1,dfn_a+siz_a-1]\(中深度为\)[dep_a,dep_a+k]$的点的$siz$之和。发现有两个限制，可以主席树维护。

代码

#include<bits/stdc++.h>

using namespace std;
const int N=300005;
const int M=N*21;
typedef long long LL;

inline int rd(){
	int x=0,f=1;char ch=getchar();
	while(!isdigit(ch)) f=ch=='-'?0:1,ch=getchar();
	while(isdigit(ch)) x=(x<<1)+(x<<3)+ch-'0',ch=getchar();
	return f?x:-x;
}

int n,q,head[N],cnt,to[N<<1],nxt[N<<1],siz[N],dfn[N],dep[N],num,bl[N];
LL ans;

struct Persistence_Segment_Tree{
	LL sum[M];int ls[M],rs[M],rt[N],tot;
	void build(int &x,int l,int r){
		x=++tot; if(l==r) return; int mid=(l+r)>>1;
		build(ls[x],l,mid); build(rs[x],mid+1,r);
	}
	void update(int pre,int &x,int l,int r,int pos,int k){
		x=++tot; sum[x]=sum[pre]+k; if(l==r) return;
		ls[x]=ls[pre]; rs[x]=rs[pre]; int mid=(l+r)>>1;
		if(pos<=mid) update(ls[pre],ls[x],l,mid,pos,k);
		else update(rs[pre],rs[x],mid+1,r,pos,k);
	}
	inline void BUILD(){
		build(rt[0],1,n);
		for(int i=1;i<=n;i++)
			update(rt[i-1],rt[i],1,n,dep[bl[i]],siz[bl[i]]-1);
	}
	LL query(int u,int v,int l,int r,int L,int R){
		if(L<=l && r<=R) return sum[v]-sum[u];
		int mid=(l+r)>>1; LL ret=0;
		if(L<=mid) ret+=query(ls[u],ls[v],l,mid,L,R);
		if(mid<R) ret+=query(rs[u],rs[v],mid+1,r,L,R);
		return ret;
	}
}tree;

inline void add(int bg,int ed){
	to[++cnt]=ed,nxt[cnt]=head[bg],head[bg]=cnt;
}

void dfs(int x,int F,int d){
	dep[x]=d; siz[x]=1; dfn[x]=++num; bl[num]=x;
	for(int i=head[x];i;i=nxt[i]){
		int u=to[i]; if(u==F) continue;
		dfs(u,x,d+1); siz[x]+=siz[u];
	}
}

int main(){
	n=rd(),q=rd(); int x,y;
 	for(int i=1;i<n;i++){
		x=rd(),y=rd();
		add(x,y),add(y,x);
	}
	dfs(1,0,1); tree.BUILD();
	while(q--){
		x=rd(),y=rd(); ans=(LL)min(dep[x]-1,y)*(siz[x]-1);
		ans+=tree.query(tree.rt[dfn[x]],tree.rt[dfn[x]+siz[x]-1],1,n,dep[x],min(n,dep[x]+y));
		printf("%lld\n",ans);
	}
	return 0;
}