bzoj 4771: 七彩树

Description

给定一棵n个点的有根树，编号依次为1到n，其中1号点是根节点。每个节点都被染上了某一种颜色，其中第i个节
点的颜色为c[i]。如果c[i]=c[j]，那么我们认为点i和点j拥有相同的颜色。定义depth[i]为i节点与根节点的距离
，为了方便起见，你可以认为树上相邻的两个点之间的距离为1。站在这棵色彩斑斓的树前面，你将面临m个问题。
每个问题包含两个整数x和d，表示询问x子树里且depth不超过depth[x]+d的所有点中出现了多少种本质不同的颜色
。请写一个程序，快速回答这些询问。

Solution

如果没有深度限制,这就是个经典的容斥,直接按照 \(dfs\) 序做
考虑相同颜色的,因为在同一子树内只算一次,我们直接在 \(dfs\) 序相邻的 \(lca\) 上 \(-1\) 就好了
答案就是一个子树权值和,线段树维护一下就好了
如果有深度限制,我们离线询问按照深度做,按深度从小到大加入, \(set\) 维护每一种颜色 \(dfs\) 序相邻的
如果强制在线的话,我们把每一个深度的线段树都存下来就好了,主席树维护一下即可,相当于转离线了

#include<bits/stdc++.h>
using namespace std;
template<class T>void gi(T &x){
	int f;char c;
	for(f=1,c=getchar();c<'0'||c>'9';c=getchar())if(c=='-')f=-1;
	for(x=0;c<='9'&&c>='0';c=getchar())x=x*10+(c&15);x*=f;
}
const int N=1e5+10;
int n,m,C[N],fa[N][18],head[N],nxt[N*2],to[N*2],num=0;
int dep[N],dfn[N],b[N],DFN=0,R[N],p[N],mx=0,rt[N],tt=0;
inline void link(int x,int y){nxt[++num]=head[x];to[num]=y;head[x]=num;}
inline void dfs(int x){
	b[dfn[x]=++DFN]=x;
	for(int i=1;i<=17;i++)fa[x][i]=fa[fa[x][i-1]][i-1];
	for(int i=head[x];i;i=nxt[i])dep[to[i]]=dep[x]+1,dfs(to[i]);
	R[x]=DFN;
}
inline int lca(int x,int y){
	if(dep[x]<dep[y])swap(x,y);
	for(int i=17;i>=0;i--)if((dep[x]-dep[y])>>i&1)x=fa[x][i];
	if(x==y)return x;
	for(int i=17;i>=0;i--)if(fa[x][i]!=fa[y][i])x=fa[x][i],y=fa[y][i];
	return fa[x][0];
}
struct data{int ls,rs,w;}tr[N*60];
inline void ins(int &x,int l,int r,int sa,int t){
	tr[++tt]=tr[x];x=tt;tr[x].w+=t;
	if(l==r)return ;
	int mid=(l+r)>>1;
	if(sa<=mid)ins(tr[x].ls,l,mid,sa,t);
	else ins(tr[x].rs,mid+1,r,sa,t);
}
inline int qry(int x,int l,int r,int sa,int se){
	if(sa<=l && r<=se)return tr[x].w;
	int mid=(l+r)>>1;
	if(se<=mid)return qry(tr[x].ls,l,mid,sa,se);
	if(sa>mid)return qry(tr[x].rs,mid+1,r,sa,se);
	return qry(tr[x].ls,l,mid,sa,mid)+qry(tr[x].rs,mid+1,r,mid+1,se);
}
inline bool comp(const int &i,const int &j){return dep[i]<dep[j];}
set<int>S[N];set<int>::iterator it1,it2,it;
inline void solve(int x){
	int c=C[x],d=dep[x];
	ins(rt[d],1,n,dfn[x],1);
	it1=S[c].lower_bound(dfn[x]);it2=S[c].upper_bound(dfn[x]);
	if(it1!=S[c].begin())it=it1,--it,ins(rt[d],1,n,dfn[lca(x,b[*it])],-1);
	if(it2!=S[c].end())it=it2,ins(rt[d],1,n,dfn[lca(x,b[*it])],-1);
	if(it1!=S[c].begin() && it2!=S[c].end()){
		--it1;
		ins(rt[d],1,n,dfn[lca(b[*it1],b[*it2])],1);
	}
	S[c].insert(dfn[x]);
}
inline void Clear(){
	DFN=num=tt=0;
	for(int i=0;i<N;i++)head[i]=rt[i]=0,set<int>().swap(S[i]);
}
inline void work(){
	Clear();
	cin>>n>>m;
	for(int i=1;i<=n;i++)gi(C[i]),p[i]=i;
	for(int i=2;i<=n;i++)gi(fa[i][0]),link(fa[i][0],i);
	dep[1]=1;dfs(1);
	sort(p+1,p+n+1,comp);mx=dep[p[n]];
	for(int i=1;i<=n;i++){
		int j=i;
		while(j<n && dep[p[j+1]]==dep[p[i]])j++;
		rt[dep[p[i]]]=rt[dep[p[i]]-1];
		for(int k=i;k<=j;k++)solve(p[k]);
		i=j;
	}
	int x,y,ans=0;
	while(m--){
		gi(x);gi(y);x^=ans;y^=ans;
		printf("%d\n",ans=qry(rt[min(mx,dep[x]+y)],1,n,dfn[x],R[x]));
	}
}
int main(){
	freopen("pp.in","r",stdin);
	freopen("pp.out","w",stdout);
	int T;cin>>T;
	while(T--)work();
	return 0;
}