이 제출은 이전 버전의 oj.uz에서 채점하였습니다. 현재는 제출 당시와는 다른 서버에서 채점을 하기 때문에, 다시 제출하면 결과가 달라질 수도 있습니다.
#include "split.h"
#include <iostream>
#include <queue>
#include <bits/stdc++.h>
using namespace std;
vector<vector<int>> adj;
int s1,s2,c1,c2,c3,n;
bool done=false;
vector<int> res;
int sub[100002];
int sz[5];
void calc(int u, int p){
	for(auto v:adj[u]){
		if(v==p) continue;
		calc(v,u);
		sub[u] += sub[v];
	}
	sub[u]++;
}
void color(int u, int p, int c, int size){
	if(size==sz[c]) return;
	queue<pair<int,int> > q;
	q.push({u,p});
	while(!q.empty() && size>sz[c]){
	    int v=q.front().first;
	    int par = q.front().second;
	    q.pop();
        res[v] = c;
        sz[c]++;
        for(auto vv:adj[v]){
            if(vv == par) continue;
            q.push({vv,v});
        }
	}
}
void solve(int u, int p){
    //printf("here\n");
	if(done) return;
	int maxi=0,idx=-1;
	for(auto v:adj[u]){
		if(sub[v]>=s1 && n-sub[v]>=s2){
			color(v,u,c1,s1);
			color(u,v,c2,s2);
			for(int i=0;i<n;i++) if(res[i]==0) res[i]=c3;
			done=true;
			return;
		}
		if(sub[v]>=s2 && n-sub[v]>=s1){
			color(v,u,c2,s2);
			color(u,v,c1,s1);
			for(int i=0;i<n;i++) if(res[i]==0) res[i]=c3;
			done=true;
			return;
		}
		if(v==p) continue;
		if (sub[v]>maxi){
			maxi = sub[v];
			idx = v;
		}
	}
	
	if(idx!=-1){
		sub[u] = n-sub[idx];
		solve(idx,u);
	}
}
vector<int> find_split(int N, int a, int b, int c, vector<int> p, vector<int> q) {
	n=N;
	adj.resize(n+2,vector<int>());
	for(int i=0; i<(int)p.size();i++){
		adj[q[i]].push_back(p[i]);
		adj[p[i]].push_back(q[i]);
	}
	res.resize(n);
	for(int i=0;i<n;i++) res[i]=0;
	vector<pair<int,int> > s;
	s.push_back({a,1});
	s.push_back({b,2});
	s.push_back({c,3});
	sort(s.begin(),s.end());
	s1 = s[0].first;
	s2 = s[1].first;
	c1 = s[0].second;
	c2 = s[1].second;
	c3 = s[2].second;
	calc(0,-1);
	solve(0,-1);
	return res;
}
| # | Verdict | Execution time | Memory | Grader output | 
|---|
| Fetching results... | 
| # | Verdict | Execution time | Memory | Grader output | 
|---|
| Fetching results... | 
| # | Verdict | Execution time | Memory | Grader output | 
|---|
| Fetching results... | 
| # | Verdict | Execution time | Memory | Grader output | 
|---|
| Fetching results... | 
| # | Verdict | Execution time | Memory | Grader output | 
|---|
| Fetching results... |