Regions 题解

Regions

这里提供一种时间复杂度不那么优秀但十分好写也好理解的做法。

题目大意

给定一颗 \(n\) 个节点的树，每个节点拥有一个颜色，进行若干次询问，每次询问给出两种颜色 \(A,B\)，求所有颜色为 \(A\) 的节点的子树中颜色为 \(B\) 的节点的个数的和。

思路分析

考虑根号分治。按颜色的节点数进行分类，节点数 \(>\sqrt n\) 的称为重颜色，节点数 \(\le \sqrt n\) 的称为轻颜色。

对询问进行分类讨论：

• \(A,B\) 均为轻颜色：

考虑到两者的节点数都不会很多，可以暴力标记每一个颜色为 \(B\) 的点，并枚举所有颜色为 \(A\) 的点，求出其子树中的标记数并累加入答案。这一过程可以用树状数组实现。单次实现的时间复杂度为 \(O(\sqrt n\log n)\)。

• \(A\) 为重颜色，\(B\) 为轻颜色：

设 \(f_{i,j}\) 表示点 \(i\) 到根的路径中颜色为重颜色 \(j\) 的节点的数量，那么答案显然为：

\[\sum\limits_{i\in \text{col}_B}f_{i,A}
\]

其中，\(\text{col}_X\) 表示所有颜色为 \(X\) 的点所构成的集合。

考虑计算 \(f\)：

枚举每一个重颜色 \(K\)，遍历 \(\text{col}_K\)，将其中的每一个点的子树中的所有点的权值加 \(1\)，再遍历每一个颜色不为 \(K\) 的点，其权值就是其对应的 \(f\)。这一过程可以用差分实现。时间复杂度为 \(O(n\sqrt n)\)。

• \(B\) 为重颜色，\(A\) 为轻颜色：

设 \(g_{i,j}\) 表示点 \(i\) 的子树中颜色为重颜色 \(j\) 的节点的数量，那么答案显然为：

\[\sum_{i\in \text{col}_A}g_{i,B}
\]

\(g\) 的计算与 \(f\) 正好相反:

枚举每一个重颜色 \(K\)，遍历 \(\text{col}_K\)，将其中的每一个点的权值单点加 \(1\)。再遍历每一个颜色不为 \(K\) 的点，其子树中的权值和即为其对应的 \(g\)，这一过程可以用前缀和实现。时间复杂度为 \(O(n\sqrt n)\)。

• \(A,B\) 均为重颜色：

容易发现，这种情况是上述两种情况的交集，上述两种情况的答案计算方式均适用于此情况。

实现细节

在实现中，为了优化空间复杂度，可以省去 \(f\) 和 \(g\) 直接累加答案。具体的说，设 \(\text{ans1}_{i,j}\) 表示所有颜色为 \(i\) 的点的子树中颜色为重颜色 \(j\) 的节点的数量的和，\(\text{ans2}_{i,j}\) 表示所有颜色为 \(i\) 的点到根的路径中颜色为重颜色 \(j\) 的节点的数量的和。（也就是第三，二种情况对应的答案）这样可以将空间复杂度由 \(O(n\sqrt n)\) 优化到 \(O(R\sqrt n)\)。

差分和前缀和可以直接在 dfs 序上进行，不需要放到树上。

总时间复杂度为 \(O(q\sqrt n\log n+n\sqrt n)\)，瓶颈在于树状数组，卡常后可以通过。

空间复杂度为 \(O(R\sqrt n+n)\)，完全可过。

代码

#include <iostream>
#include <cstring>
#include <cstdio>
#include <cmath>
#include <algorithm>
#include <vector>

using namespace std;
const int N=200100,M=450,R=25500;

int idx=1,n,m,q,B,cnt,in1,in2,cur,ans;
int w[N],to[N],nxt[N],head[N];
int c[N],dfn[N],siz[N],big[N],d[N];
int ans1[M][R],ans2[M][R];

vector<int> col[N];

int read(){//卡常用的快读快写
    int x=0;char ch=getchar();
    while(ch>'9'||ch<'0') ch=getchar();
    while(ch>='0'&&ch<='9') x=(x<<3)+(x<<1)+ch-'0',ch=getchar();
	return x;
}

void write(int x){
    int k=x/10;if(k) write(k);
    putchar(x-(k<<3)-(k<<1)+'0');
}

void add(int u,int v){
	idx++;to[idx]=v;nxt[idx]=head[u];head[u]=idx;
}

void dfs_1(int s){
	dfn[s]=++cnt;siz[s]=1;
	for(int i=head[s];i;i=nxt[i])
		dfs_1(to[i]),siz[s]+=siz[to[i]];
}

struct BIT{
	int a[N];
	#define lowbit(x) ((x)&(-(x)))
	void add(int x,int v){
		for(;x<N;x+=lowbit(x)) a[x]+=v;
	}
	int query(int x){
		if(!x) return 0;
		int ans=0;
		for(;x;x-=lowbit(x)) ans+=a[x];
		return ans;
	}
}tree;

int main(){
	n=read();m=read();q=read();
	B=sqrt(n);
	w[1]=read();
	col[w[1]].push_back(1);
	for(int i=2;i<=n;i++){
		in1=read();w[i]=read();
		add(in1,i);
		col[w[i]].push_back(i);
	}
	dfs_1(1);
	for(int i=1;i<=m;i++){
		if(col[i].size()<=B) continue;
		big[i]=++cur;//对每一种重颜色重标号
		for(int j=1;j<=n;j++) d[j]=0;
		for(auto it:col[i]) d[dfn[it]]++;//单点加
		for(int j=1;j<=n;j++) d[j]+=d[j-1];//前缀和
		for(int j=1;j<=n;j++){
			if(w[j]==i) continue;
			ans1[cur][w[j]]+=d[dfn[j]+siz[j]-1]-d[dfn[j]-1];//直接累加答案
		}
		for(int j=1;j<=n;j++) d[j]=0;
		for(auto it:col[i]){
			d[dfn[it]]++;//差分
			d[dfn[it]+siz[it]]--;//这里其实是 dfn[it]+siz[it]-1+1
		}
		for(int j=1;j<=n;j++) d[j]+=d[j-1];//对差分序列做前缀和得到原序列
		for(int j=1;j<=n;j++){
			if(w[j]==i) continue;
			ans2[cur][w[j]]+=d[dfn[j]];
		}
	}
	while(q--){
		in1=read();in2=read();ans=0;
		if(big[in2]){write(ans1[big[in2]][in1]);cout<<'\n'<<flush;continue;}
		if(big[in1]){write(ans2[big[in1]][in2]);cout<<'\n'<<flush;continue;}
		for(auto it:col[in2]) tree.add(dfn[it],1);//暴力加
		for(auto it:col[in1])
			ans+=tree.query(dfn[it]+siz[it]-1)-tree.query(dfn[it]-1);
		write(ans);cout<<'\n'<<flush;
		for(auto it:col[in2]) tree.add(dfn[it],-1);//记得清空
	}
	return 0;
}