ABC202E Count Descendants

ABC202E Count Descendants

线段树合并模板题。

每次询问就是给定有序数对 \((u,d)\)，求有根树 \(T\) 上，点 \(u\) 的子树内有多少点 \(v\)，使得 \(v\) 的深度恰好等于 \(d+1\)。定义根节点深度为 \(1\)。

考虑对每一个点开一个长度为 \(n\) （因为 \(T\) 的最大深度为 \(n\)）的数组 \(a_u\)，\(a_{u,i}\) 表示 \(u\) 的子树内深度为 \(i\) 的点有多少，同时记录每个点的深度 \(depth\)。

每递归到叶节点 \(v\)，就将 \(a_{v,i}\) 加上 \(1\)。

回溯时，将 \(a_v\) 暴力合并到 \(a_u\)，即 \(\forall i \in [1,n],a_{u,i} \leftarrow a_{u,i}+a_{v,i}\)，这一步的时间复杂度为 \(\mathcal{O}(n)\)。单次查询可以做到 \(\mathcal{O}(1)\)。于是该做法的时间、空间复杂度均为 \(\mathcal{O}(n^2)\)。

显然，时间复杂度的瓶颈在于合并操作，考虑如何优化这一步。

考虑线段树合并，对每一个 \(u\) 开一棵动态开点权值线段树，每次回溯时将 \(v\) 对应的线段树与 \(u\) 合并即可，单次时间复杂度 \(\mathcal{O}( \log n)\)。

注意不能将 \(v\) 对应的线段树直接合并到 \(u\)，因为这样会导致 \(u\) 对应线段树的信息丢失，所以应该先新建一个节点 \(o\)，再将 \(u\) 和 \(v\) 的节点信息合并至 \(o\)。当然离线询问也可以解决这个问题。

由于每次会多开 \(\mathcal{O}(\log n)\) 个点，所以线段树的数组大小要在原来的基础上再多开一倍。总复杂度 \(\mathcal{O}(n \log n)\)。

树上启发式合并 也可以解决本题，时间复杂度仍然是 \(\mathcal{O}(n \log n)\)。

下面是线段树合并的代码：

#include<bits/stdc++.h>
using namespace std;
bool Mbegin;
void File_Work(){
	freopen("test.in","r",stdin);
	freopen("test.out","w",stdout);
}
namespace Fast_In_Out{
	char buf[1<<21],*P1=buf,*P2=buf;
	#define gc (P1==P2&&(P2=(P1=buf)+fread(buf,1,1<<21,stdin),P1==P2)?EOF:*P1++)
	int read(){
		int f=1,x=0;
		char c=gc;
		while(c<'0'||c>'9'){
			if(c=='-')
			    f=-1;
			c=gc;
		}
		while(c>='0'&&c<='9'){
			x=x*10+c-'0';
			c=gc;
		}
		return f*x;
	}
	void write(int x){
		if(x<0)
			x=-x,putchar('-');
		if(x>9)
			write(x/10);
		putchar(x%10+'0');
	}
	#undef gc
}
using namespace Fast_In_Out;
const int N=2e5+8;
int n,m;
struct Graph{
	int head[N],edge_tot=1,to[N],next[N];
	void add_edge(int u,int v){
		edge_tot++;
		to[edge_tot]=v;
		next[edge_tot]=head[u];
		head[u]=edge_tot;
	}
}Tree;
int depth[N];
int rt[N];
struct Segemnt_Tree{
	int ocnt,ls[N<<6],rs[N<<6],sum[N<<6];
	void push_up(int o){
		sum[o]=sum[ls[o]]+sum[rs[o]];
	}
	void insert(int &o,int l,int r,int pos){
		o=++ocnt;
		if(l==r){
			sum[o]++;
			return;
		}
		int mid=(l+r)/2;
		if(pos<=mid)
			insert(ls[o],l,mid,pos);
		else
			insert(rs[o],mid+1,r,pos);
		push_up(o);
	}
	int query(int o,int l,int r,int pos){
		if(o==0)
			return 0;
		if(l==r)
			return sum[o];
		int mid=(l+r)/2;
		if(pos<=mid)
			return query(ls[o],l,mid,pos);
		else
			return query(rs[o],mid+1,r,pos);
	}
	int merge(int u,int v){
		if(u==0||v==0)
			return u+v;
		int o=++ocnt;
		ls[o]=merge(ls[u],ls[v]);
		rs[o]=merge(rs[u],rs[v]);
		sum[o]=sum[u]+sum[v];
		return o;
	}
}smt;
void dfs(int u,int father){
	depth[u]=depth[father]+1;
	smt.insert(rt[u],1,n,depth[u]);
	for(int i=Tree.head[u];i;i=Tree.next[i]){
		int v=Tree.to[i];
		dfs(v,u);
		rt[u]=smt.merge(rt[u],rt[v]);
	}
}
bool Mend;
int main(){
//	File_Work();
	fprintf(stderr,"%.3lf MB\n\n\n",(&Mbegin-&Mend)/1048576.0);
	n=read();
	for(int i=2;i<=n;i++){
		int father=read();
		Tree.add_edge(father,i);
	}
	dfs(1,0);
	m=read();
	while(m--){
		int a=read(),b=read()+1;
		write(smt.query(rt[a],1,n,b)),putchar('\n');
	}
	fprintf(stderr,"\n\n\n%.0lf ms",1e3*clock()/CLOCKS_PER_SEC);
	return 0;
}

线段树合并的题，通常可以先想 \(\mathcal{O}(n^2)\) 做法，再进行优化。

思想几乎一样的题：CF208E Blood Cousins。

另一道线段树合并模板题：CF600E Lomsat gelral。