IOI 2009

regions.cpp

// Ivan Carvalho
// Regions - IOI 2009
// O(N*sqrt(Q))
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<int,int> ii;
const int MAXN = 2*1e5 + 10;

//Declarando as variáveis
vector<int> grafo[MAXN],regioes[MAXN];
vector<ii> perguntas[MAXN];
int ini[MAXN],fim[MAXN],dfsPtr,cor[MAXN],freq[MAXN],pesado[MAXN],sz[MAXN],vetor[MAXN],N,R,Q;
ll tot[MAXN],exibe[MAXN];

// Utilidades para o Sack
void add(int v){
	freq[cor[v]]++;
}
void remove(int v){
	freq[cor[v]]--;
}
// Calculando o tamanho das sub-árvores

int dfs1(int v){
	ini[v] = ++dfsPtr;
	vetor[ini[v]] = v;
	sz[v] = 1;
	for(int u : grafo[v]){
		sz[v] += dfs1(u);
	}
	fim[v] = dfsPtr;
	return sz[v];
}

// Rodando o Sack
void dfs2(int v,int keep){
	// Big guarda o filho de maior tamanho
	int big = -1,mx = 0;
	for(int u : grafo[v]){
		if(sz[u] > mx){
			big = u;
			mx = sz[u];
		}
	}
	// Não salvamos as informações dos filhos que não são o maior
	for(int u : grafo[v]){
		if(u != big) dfs2(u,0);
	}
	// Caso exista o maior, salvamos sua informação
	if(big != -1) dfs2(big,1);
	add(v);
	// Adicionamos as informações dos filhos menores
	for(int u : grafo[v]){
		if(u == big) continue;
		for(int i = ini[u];i<=fim[u];i++){
			add(vetor[i]);
		}
	}
	// Caso não haja muitas perguntas na região, ajudamos a responder-las agora 
	if(!pesado[cor[v]]){
		for(ii par : perguntas[cor[v]]){
			exibe[par.second] += freq[par.first];
		}
	}
	// Se não devemos guardar a informação desse vértice, limpamos
	if(keep == 0){
		for(int i = ini[v];i<=fim[v];i++){
			remove(vetor[i]);
		}
	}
}

// No caso das regiões com muitas queries, processamos suas perguntas separadamente
void processa_heavy(int familia){
	// Limpamos os arrays
	memset(freq,0,sizeof(freq));
	memset(tot,0,sizeof(tot));
	// Fazemos a soma acumulada no vetor da árvore
	for(int v : regioes[familia]){
		freq[ini[v]]++;
		freq[fim[v]+1]--;
	}
	for(int i = 1;i<=N;i++){
		freq[i] += freq[i-1];
	}
	// Adicionamos quantos ancestrais de i são dessa região
	for(int i = 1;i<=N;i++){
		tot[cor[i]] += 1LL*freq[ini[i]];
	}
	// Respondemos as queries
	for(ii par : perguntas[familia]){
		exibe[par.second] = tot[par.first];
	}
}

int main(){
	// Leemos a entrada
	scanf("%d %d %d",&N,&R,&Q);
	scanf("%d",&cor[1]);
	regioes[cor[1]].push_back(1);
	for(int i = 2;i<=N;i++){
		int s;
		scanf("%d %d",&s,&cor[i]);
		grafo[s].push_back(i);
		regioes[cor[i]].push_back(i);
	}
	// Rodamos o primeiro dfs e lemos a entrada
	dfs1(1);
	for(int i = 1;i<=Q;i++){
		int r1,r2;
		scanf("%d %d",&r1,&r2);
		perguntas[r1].push_back(ii(r2,i));
	}
	// Sqrt-decomposition e Sack
	int bucket = (int)ceil(sqrt(Q));
	for(int i = 1;i<=R;i++){
		if(perguntas[i].size() >= bucket) pesado[i] = 1;
	}
	dfs2(1,0);
	for(int i = 1;i<=R;i++){
		if(pesado[i]) processa_heavy(i);
	}
	// Exibimos a resposta
	for(int i = 1;i<=Q;i++) printf("%lld\n",exibe[i]);
	return 0;
}