【BZOJ3351】Regions（IOI2009）-分塊+vector

測試地址：Regions
題目大意： 給定一棵$n$個點的有根樹，每個點有顏色，$q$個詢問，每次詢問給出兩個顏色$a,b$，表示詢問樹中有多少對點$(u,v)$使得$u$顏色是$a$，$v$顏色是$b$，且$u$是$v$的祖先。
做法： 本題需要用到分塊+vector。
好題。傳統數據結構貌似做不了這道題，於是想到分塊。
很容易想到兩種暴力：
1.將$a$相同的詢問一起處理，對所有點統計合法的祖先數目計算。
2.將$b$相同的詢問一起處理，對所有點統計合法的子孫數目計算。
上面兩種算法的總時間複雜度都是$O(rn)$的，無法通過此題。但我們想到一個常用的分塊思路：設定閾值分類討論。於是在這裏，我們對$b$這種顏色的出現次數$B$進行討論。
當$B>\sqrt n$時，滿足這種條件的$b$只有$\sqrt n$種，於是採用上面第二種暴力，時間複雜度$O(n\sqrt n)$。
當$B\le \sqrt n$時，每個詢問涉及到的$b$顏色的點最多有$\sqrt n$個，所以我們可以在每種顏色上掛一個vector（在每個點上掛空間會爆炸），存儲$b$爲這種顏色的詢問，然後一次DFS，DFS的同時維護從根到當前點的路徑上各顏色點出現的次數，當走到一個$B\le \sqrt n$的點時，對這種顏色所影響的詢問進行更新。我們發現這差不多就是第一種暴力的思路，只不過我們同時處理了所有的$a$。因爲有$q$個詢問，每個詢問被$\sqrt n$個點影響，所以時間複雜度爲$O(q\sqrt n)$。
於是我們就解決了這一題，時間複雜度爲$O((n+q)\sqrt n)$。
以下是本人代碼：

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
int n,R,q,r[200010],first[200010]={0},tot=0,blocksiz;
int opx[200010],opy[200010];
ll ans[200010]={0},cnt[200010]={0},sum[25010]={0};
vector<int> pos[25010],posq[25010];
struct edge
{
	int v,next;
}e[200010];

void insert(int a,int b)
{
	e[++tot].v=b;
	e[tot].next=first[a];
	first[a]=tot;
}

void dfs1(int v)
{
	int siz=posq[r[v]].size();
	if (pos[r[v]].size()<=blocksiz)
	{
		for(int i=0;i<siz;i++)
			ans[posq[r[v]][i]]+=cnt[opx[posq[r[v]][i]]];
	}
	cnt[r[v]]++;
	for(int i=first[v];i;i=e[i].next)
		dfs1(e[i].v);
	cnt[r[v]]--;
}

void dfs2(int v)
{
	ll now=cnt[v];
	for(int i=first[v];i;i=e[i].next)
	{
		dfs2(e[i].v);
		cnt[v]+=cnt[e[i].v];
	}
	sum[r[v]]+=cnt[v]-now;
}

int main()
{
	scanf("%d%d%d",&n,&R,&q);
	
	scanf("%d",&r[1]);
	pos[r[1]].push_back(1);
	for(int i=2;i<=n;i++)
	{
		int fa;
		scanf("%d%d",&fa,&r[i]);
		pos[r[i]].push_back(i);
		insert(fa,i);
	}
	
	blocksiz=(int)sqrt(n+1);
	for(int i=1;i<=q;i++)
	{
		scanf("%d%d",&opx[i],&opy[i]);
		posq[opy[i]].push_back(i);
	}
	dfs1(1);
	
	for(int i=1;i<=R;i++)
		if (pos[i].size()>blocksiz)
		{
			memset(cnt,0,sizeof(cnt));
			memset(sum,0,sizeof(sum));
			int siz=pos[i].size();
			for(int j=0;j<siz;j++)
				cnt[pos[i][j]]++;
			dfs2(1);
			siz=posq[i].size();
			for(int j=0;j<siz;j++)
				ans[posq[i][j]]=sum[opx[posq[i][j]]];
		}
	
	for(int i=1;i<=q;i++)
		printf("%lld\n",ans[i]);
	
	return 0; 
}