Submission #96627

#	Time	Username	Problem	Language	Result	Execution time	Memory
96627		KLPP	Balanced Tree (info1cup18_balancedtree)	C++14	0 / 100	1762 ms	525312 KiB

balancedtree

#include<bits/stdc++.h>

using namespace std;
typedef long long int lld;
vector<int>nei[1000000];
int n;
int color[1000000];
int dist[1000000];
class DSU{
	int parent[1000000];
	int size[1000000];
	public:
	int isolated;
	void init(int n,int C){isolated=0;
		for(int i=0;i<n;i++){
			if(color[i]==C)size[i]=1,isolated++;
			else size[i]=0;
		}
		for(int i=n;i<2*n-1;i++)size[i]=0;
		for(int i=0;i<2*n-1;i++)parent[i]=i;
	}
	int root(int x){
		if(parent[x]==x)return x;
		parent[x]=root(parent[x]);
		return parent[x];
	}
	void merge(int a, int b){
		a=root(a);
		b=root(b);
		//cout<<a<<" "<<b<<endl;
		if(a==b)return;
		if(size[a]==1 && size[b]!=0)isolated--;
		if(size[b]==1 && size[a]!=0)isolated--;
		if(size[a]<size[b])swap(a,b);
		size[a]+=size[b];
		parent[b]=a;
	}
	void print(int n){
		for(int i=0;i<2*n-1;i++)cout<<parent[i]<<" "<<size[i]<<endl;
	}
};
int solve(int C){
	DSU *D=new DSU();
	queue<int> q;
	for(int i=0;i<2*n-1;i++){
		dist[i]=-1;
		if(i<n && color[i]==C){//cout<<i<<endl;
			dist[i]=0;
			q.push(i);
		}
	}
	D->init(n,C);
	
	while(!q.empty() && D->isolated>0){
		int u=q.front();q.pop();
		for(int i=0;i<nei[u].size();i++){
			int v=nei[u][i];
			D->merge(u,v);
			if(dist[v]==-1){
				dist[v]=dist[u]+1;
				q.push(v);
			}
		}
	}
	//D->print(n);
	if(D->isolated>0)return 1000000000;
	int ans=0;
	for(int i=0;i<2*n-1;i++){
		ans=max(ans,dist[i]);
		//cout<<dist[i]<<endl;
	}
	return ans;
}
int main(){
	int t;
	cin>>t;
	while(t--){
		cin>>n;
		for(int i=0;i<2*n-1;i++)nei[i].clear();
		for(int i=0;i<n-1;i++){
			int x,y;
			cin>>x>>y;
			x--;y--;
			nei[x].push_back(i+n);
			nei[i+n].push_back(x);
			nei[y].push_back(i+n);
			nei[i+n].push_back(y);			
		}
		for(int i=0;i<n;i++)cin>>color[i];
		int D=0;
		D=max(D,solve(0));
		D=max(D,solve(1));
		if(D>n)cout<<-1<<endl;
		else cout<<D<<endl;
	}
	return 0;
}

Compilation message (stderr)

balancedtree.cpp: In function 'int solve(int)':
balancedtree.cpp:56:16: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
   for(int i=0;i<nei[u].size();i++){
               ~^~~~~~~~~~~~~~

• Subtask 1: 0 / 10
• Subtask 2: 0 / 13
• Subtask 3: 0 / 26
• Subtask 4: 0 / 45
• Subtask 5: 0 / 6