Cod sursa(job #2147430)

Utilizator flibiaVisanu Cristian flibia Data 28 februarie 2018 18:58:08
Problema Cerere Scor 80
Compilator cpp Status done
Runda Arhiva de probleme Marime 1.91 kb
#include <bits/stdc++.h>

using namespace std;

#define dim 100000
char buff[dim];
int p = 0;

void read(int &nr){
	nr = 0;
	while(buff[p] < '0' || buff[p] > '9')
		if(++p == dim) 
			fread(buff, 1, dim, stdin), p = 0;
	while(buff[p] >= '0' && buff[p] <= '9'){
		nr = 10 * nr + buff[p] - '0';
		if(++p == dim) 
			fread(buff, 1, dim, stdin), p = 0;
	}
}

int n, nr[100100], dp[100100], pv[100100], anc[20][100100], x, y, idx;
bool bad[100100], viz[100100];
vector <int> level[100100], v[100100];

void build(){
	int sz = log2(n);
	for(int i = 1; i <= n; i++)
		anc[0][i] = i;
	for(int lvl = 1; lvl <= sz; lvl++)
		for(int i = 1; i <= n; i++)
			x = anc[lvl - 1][i], anc[lvl][i] = anc[lvl - 1][pv[x]];
}

int solve(int nod, int p){
	if(p == 0)
		return nod;
	int lvl = 0;
	while((1 << lvl) - 1 <= p)
		lvl++;
	lvl--;
	if(anc[lvl][nod] == 0)
		return 0;
	return solve(anc[lvl][nod], p - (1 << lvl) + 1); 
}

int dfs(int nod, int lvl){
	if(dp[nod] != -1 && lvl == 0)
		return 1 + dp[nod];
	if(lvl == 0 && dp[nod] == -1)
		return 1 + (dp[nod] = dfs(nod, nr[nod]));
	else return dfs(solve(nod, lvl), 0);
}

void predfs(int from, int lvl){
	level[lvl].push_back(from);
	viz[from] = 1;
	for(auto to : v[from])
		if(!viz[to])
			predfs(to, lvl + 1);
}

int main(){
	freopen("cerere.in", "r", stdin);
	freopen("cerere.out", "w", stdout);
	memset(dp, -1, sizeof dp);
	cin >> n;
	for(int i = 1; i <= n; i++){
		read(nr[i]);
		if(!nr[i])
			dp[i] = 0;
	}
	for(int i = 1; i < n; i++){
		read(x);
		read(y);
		pv[y] = x;
		bad[y] = 1;
		v[x].push_back(y);
	}
	for(int i = 1; i <= n; i++)
		if(!bad[i])
			idx = i;
	predfs(idx, 1);
	build();
//	for(int i = 1; i <= n; i++)
//		if(nr[i])
//			dp[i] = dfs(i, nr[i]);
	for(int i = 1; i <= n; i++)
		for(auto j : level[i])
			if(nr[j])
				dp[j] = dfs(j, nr[j]);
	for(int i = 1; i <= n; i++)
		cout << dp[i] << ' ';
	return 0;
}