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$，即 $\mathrm{\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)$。

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

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。