P3899 [湖南集训]谈笑风生 - 线段树合并

题解

对于每个点，以深度为下标记录答案。

如果线段树合并想在线，Merge 函数需要这样写：

int Merge(int p,int q){
	if(!p||!q) return p^q;
	int res=Newnode(t[p].l,t[p].r,t[p].v+t[q].v);
	ls(res)=Merge(ls(p),ls(q)),rs(res)=Merge(rs(p),rs(q));
	return res;
}

当 \(p\) 或 \(q\) 为空时，直接返回另一个。这保证了线段树合并的时空复杂度，但也使得修改父亲节点的线段树时可能修改到子节点。

两种解决方案：

• 将修改放到合并前面；
• 将修改可持久化。这样，每次修改一定不会影响原树。合并时只合并最后一个版本即可。

代码

#include <cstdio>
#include <cstring>
#include <cctype>
#include <vector>
using namespace std;
#define For(Ti,Ta,Tb) for(int Ti=(Ta);Ti<=(Tb);++Ti)
#define Dec(Ti,Ta,Tb) for(int Ti=(Ta);Ti>=(Tb);--Ti)
template<typename T> void Read(T &x){
	x=0;int _f=1;
	char ch=getchar();
	while(!isdigit(ch)) _f=(ch=='-'?-1:_f),ch=getchar();
	while(isdigit(ch)) x=x*10+(ch^48),ch=getchar();
	x=x*_f;
}
template<typename T,typename... Args> void Read(T &x,Args& ...others){
	Read(x);Read(others...);
}
typedef long long ll;
const int Inf=0x3f3f3f3f,N=3e5+5;
int n,q;
vector<int> G[N];
int dep[N],siz[N];
void Dfs(int u,int fa){
	dep[u]=dep[fa]+1,siz[u]=1;
	for(int v:G[u]){
		if(v==fa) continue;
		Dfs(v,u);siz[u]+=siz[v];
	}
}
#define ls(xx) t[xx].ls
#define rs(xx) t[xx].rs
struct Node{
	int l,r,ls,rs;ll v;
}t[20000000];
int tot=0;
int Newnode(int l,int r,ll v=0){
	t[++tot]=Node{l,r,0,0,v};
	return tot;
}
void Pushup(int p){t[p].v=t[ls(p)].v+t[rs(p)].v;}
void Add(int p,int pos,ll x){
	if(t[p].l==t[p].r){
		t[p].v+=x;return;
	}
	int mid=(t[p].l+t[p].r)>>1;
	if(pos<=mid){
		if(!ls(p)) ls(p)=Newnode(t[p].l,mid);
		Add(ls(p),pos,x);
	}else{
		if(!rs(p)) rs(p)=Newnode(mid+1,t[p].r);
		Add(rs(p),pos,x);
	}
	Pushup(p);
}
ll Query(int p,int l,int r){
	if(!p||r<t[p].l||l>t[p].r) return 0;
	if(l<=t[p].l&&t[p].r<=r) return t[p].v;
	return Query(ls(p),l,r)+Query(rs(p),l,r);
}
int Merge(int p,int q){
	if(!p||!q) return p^q;
	int res=Newnode(t[p].l,t[p].r,t[p].v+t[q].v);
	ls(res)=Merge(ls(p),ls(q)),rs(res)=Merge(rs(p),rs(q));
	return res;
}
int root[N];
void Work(int u,int fa){
	root[u]=Newnode(1,n);
	Add(root[u],dep[u],siz[u]-1);
	for(int v:G[u]){
		if(v==fa) continue;
		Work(v,u);
		root[u]=Merge(root[u],root[v]);
	}
}
int main(){
	Read(n,q);
	int u,v;
	For(i,1,n-1){
		Read(u,v);
		G[u].push_back(v),G[v].push_back(u);
	}
	Dfs(1,0);
	Work(1,0);
	while(q--){
		int p,k;Read(p,k);
		printf("%lld\n",1LL*min(dep[p]-1,k)*(siz[p]-1)+Query(root[p],dep[p]+1,min(dep[p]+k,n)));
	}
	return 0;
}