[BZOJ 4771]七彩树(可持久化线段树+树上差分)

[BZOJ 4771]七彩树(可持久化线段树+树上差分)

题面

给定一棵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的所有点中出现了多少种本质不同的颜色。请写一个程序，快速回答这些询问。

分析

先不考虑深度限制，我们考虑如何统计颜色。

我们把每种颜色开一个set，set里存储该颜色节点的dfs序。对于一对相同颜色的节点u,v，他们会对u到v的路径上的节点，和lca(u,v)到根节点路径上的节点的答案产生1的贡献。可以用树上差分算法。开一棵线段树，线段树第x个叶子节点存储节点dfs序（记作dfn) 为x的差分值，维护区间和。那么我们把dfn[u]+1,dfn[v]+1,dfn[lca(u,v)]-1即可。查询x子树的时候直接区间查询x的dfs序对应的区间。

那么我们加入一个新节点x的时候如何更新答案呢。这实际上是在处理树链的并。我们在set中找出x的前驱pre和后继nex（按照dfn排序后），dfn[x]+1,dfn[lca(pre,x)]-1,相当于把x到lca(pre,x)的路径加入答案。同理对nex进行操作。注意lca(pre,nex)往上的路径被减了两次，所以dfn[lca(pre,nex)] -1.这里画个图可以方便理解

那么有深度限制要怎么做呢？维护可持久化线段树，第i棵可持久化线段树存储深度<=i的树的答案，把节点按深度排序。依次加入对应深度的可持久化线段树。查询的时候直接在对应的线段树中查询x子树即可。

代码

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#include<cmath>
#include<set>
#define maxn 400000
#define maxlogn 32
using namespace std;
inline int qread(int &x){
	x=0;
	int sign=1;
	char c=getchar();
	while(c<'0'||c>'9'){
		if(c=='-') sign=-1;
		c=getchar();
	}
	while(c>='0'&&c<='9'){
		x=x*10+c-'0';
		c=getchar();
	}
	return x*sign;
}
inline void qprint(int x){
	if(x<0){
		putchar('-');
		qprint(-x);
	}else if(x==0){
		putchar('0');
		return;
	}else{
		if(x/10>0) qprint(x/10);
		putchar('0'+x%10);
	}
}
int t,n,m;
int c[maxn+5];
struct edge{
	int from;
	int to;
	int next;
}E[maxn*2+5];
int sz=1;
int head[maxn+5];
void add_edge(int u,int v){
	sz++;
	E[sz].from=u;
	E[sz].to=v;
	E[sz].next=head[u];
	head[u]=sz;
}

int log2n;
int tim;
int anc[maxn+5][maxlogn+5];
int hash_dfn[maxn+5];
int deep[maxn+5];
int lb[maxn+5],rb[maxn+5];
int id[maxn+5];
int cmp(int x,int y){
	return deep[x]<deep[y];
}

void dfs(int x,int fa){
	tim++;
	lb[x]=tim;
	deep[x]=deep[fa]+1;
	anc[x][0]=fa;
	hash_dfn[lb[x]]=x;
	for(int i=1;i<=log2n;i++) anc[x][i]=anc[anc[x][i-1]][i-1];
	for(int i=head[x];i;i=E[i].next){
		int y=E[i].to;
		if(y!=fa){
			dfs(y,x);
		}
	}
	rb[x]=tim;
}

int lca(int x,int y){
	if(deep[x]<deep[y]) swap(x,y);
	for(int i=log2n;i>=0;i--){
		if(deep[anc[x][i]]>=deep[y]){
			x=anc[x][i];
		}
	}
	if(x==y) return x;
	for(int i=log2n;i>=0;i--){
		if(anc[x][i]!=anc[y][i]){
			x=anc[x][i];
			y=anc[y][i];
		}
	}
	return anc[x][0];
}
struct per_segment_tree{
	//第i棵可持久化线段树维护deep<=i,且dfn在[l,r]内的节点的不同颜色个数 
	struct node{
		int ls;
		int rs;
		int cnt;
	}tree[maxn*maxlogn+5];
	int root[maxn+5];
	int ptr;
	void push_up(int x){
		tree[x].cnt=tree[tree[x].ls].cnt+tree[tree[x].rs].cnt;
	}
	void insert(int &x,int last,int upos,int uval,int l,int r){
		x=++ptr;
		tree[x]=tree[last];
		if(l==r){
			tree[x].cnt+=uval;
			return;
		}
		int mid=(l+r)>>1;
		if(upos<=mid) insert(tree[x].ls,tree[last].ls,upos,uval,l,mid);
		else insert(tree[x].rs,tree[last].rs,upos,uval,mid+1,r);
		push_up(x); 
	}
	int query(int x,int L,int R,int l,int r){
		if(x==0) return tree[x].cnt;
		if(L<=l&&R>=r){
			return tree[x].cnt;
		}
		int mid=(l+r)>>1;
		int ans=0;
		if(L<=mid) ans+=query(tree[x].ls,L,R,l,mid);
		if(R>mid) ans+=query(tree[x].rs,L,R,mid+1,r);
		return ans;
	}
}T;

set<int>s[maxn+5];
void ini(){
	log2n=log2(n)+1;
	tim=0;
	sz=1;
	T.ptr=0;
	for(int i=1;i<=n;i++){//不要memset,会TLE
		anc[i][0]=0;
		c[i]=0;
		deep[i]=0;
		hash_dfn[i]=0;
		head[i]=0;
		id[i]=0;
		lb[i]=0;
		rb[i]=0;
		s[i].clear();
		T.root[i]=0;
	}
}
int main(){

	int u,v;
	qread(t);
	while(t--){
		
		qread(n);
		qread(m);
		ini();
		
		for(int i=1;i<=n;i++) qread(c[i]);
		for(int i=2;i<=n;i++){
			qread(u);
			add_edge(u,i);
			add_edge(i,u);
		}
		dfs(1,0);
		for(int i=1;i<=n;i++) id[i]=i;
		sort(id+1,id+1+n,cmp);
		for(int i=1;i<=n;i++){
			int x=id[i];
			int pre=0,nex=0;

			T.insert(T.root[deep[x]],T.root[deep[id[i-1]]],lb[x],1,1,n);
			set<int>::iterator it=s[c[x]].lower_bound(lb[x]);
			if(it!=s[c[x]].begin()){
				set<int>::iterator it2=it;//不能直接--it,会影响下一个判断 
				pre=hash_dfn[*(--it2)];

				T.insert(T.root[deep[x]],T.root[deep[x]],lb[lca(pre,x)],-1,1,n);
				//x会对pre到x的路径上的点产生1的贡献，树上差分 
				//线段树里的每个店储存的都是差分值 
			}
			if(it!=s[c[x]].end()){
				nex=hash_dfn[*it];

				T.insert(T.root[deep[x]],T.root[deep[x]],lb[lca(x,nex)],-1,1,n);
			}
			//lca(pre,nex)上方的路径被多减了一次，加回来
			if(pre!=0&&nex!=0){

				T.insert(T.root[deep[x]],T.root[deep[x]],lb[lca(pre,nex)],1,1,n);
			} 
			s[c[x]].insert(lb[x]);
		}
		
		int last=0;
		int x,d;
		for(int i=1;i<=m;i++){
			qread(x);
			qread(d);
			x^=last;
			d^=last;
			last=T.query(T.root[min(deep[x]+d,deep[id[n]])],lb[x],rb[x],1,n);
			qprint(last);
			putchar('
');
		}
	}
}