# | Time | Username | Problem | Language | Result | Execution time | Memory |
---|---|---|---|---|---|---|---|
439136 | mangoesry | Split the Attractions (IOI19_split) | C++14 | 0 ms | 0 KiB |
This submission is migrated from previous version of oj.uz, which used different machine for grading. This submission may have different result if resubmitted.
#include <bits/stdc++.h>
#include "split.h"
using namespace std;
const int MX = 1e5 + 5;
vector<int> adj[MX],adjTree[MX],res,P,Q;
set<int> used;
pair<int,int> group[3];
int N,vis[MX],groupChild[MX],subTree[MX],now;
bool valid = false;
void dfsTree(int u){
vis[u] = true;
subTree[u] = 1;
for(int v : adj[u]){
if(vis[v]){
continue;
}
adjTree[u].push_back(v);
adjTree[v].push_back(u);
dfsTree(v);
subTree[u] += subTree[v];
}
}
void fillResult(int ii,int cnt,int u){
if(now == cnt || res[u]){
return;
}
now++;
res[u] = ii;
for(int v : adjTree[u]){
fillResult(ii,cnt,v);
}
}
void dfs(int u,int p){
for(int v : adjTree[u]){
if(v == p){
continue;
}
dfs(v,u);
}
if(!valid && subTree[u] >= group[0].first && N-subTree[u] >= group[0].first){
valid = true;
if(N > 2*subTree[u]){
now = 0;
res[p] = 1;
fillResult(group[0].second,group[0].first,u);
now = 0;
res[p] = 0;
fillResult(group[1].second,group[1].first,0);
}else{
now = 0;
res[p] = 1;
fillResult(group[1].second,group[1].first,u);
now = 0;
res[p] = 0;
fillResult(group[0].second,group[0].first,0);
}
}
}
void fillResultSmall(int ii,int cnt,int u){
if(now == cnt || res[u] || !used.count(u)){
return;
}
now++;
res[u] = ii;
for(int v : adj[u]){
fillResultSmall(ii,cnt,v);
}
}
void subChild(int u,int p,int vv){
groupChild[u] = vv;
for(int v : adjTree[u]){
if(v == p){
continue;
}
subChild(v,u,vv);
}
}
void dfsSmall(int u,int p){
bool small = true;
for(int v : adjTree[u]){
if(v == p || subTree[v] < group[0].first){
continue;
}
small = false;
dfsSmall(v,u);
}
if(!valid && small && subTree[u] >= group[0].first && N-subTree[u] < group[0].first){
memset(groupChild,-1,sizeof(groupChild));
for(int v : adjTree[u]){
if(v == p){
continue;
}
subChild(v,u,v);
}
groupChild[u] = u;
int sum = N-subTree[u];
set<int> choose;
choose.insert(-1);
for(int i=0;i<(int)P.size();i++){
if(sum >= group[0].first){
break;
}
if((groupChild[P[i]] == -1) ^ (groupChild[Q[i]] == -1)){
if(groupChild[P[i]] == u || groupChild[Q[i]] == u){
continue;
}
if(groupChild[P[i]] != -1 && !choose.count(groupChild[P[i]])){
sum += subTree[groupChild[P[i]]];
choose.insert(groupChild[P[i]]);
}else if(!choose.count(groupChild[Q[i]]){
sum += subTree[groupChild[Q[i]]];
choose.insert(groupChild[Q[i]]);
}
}
}
if(sum >= group[0].first){
valid = true;
used.clear();
for(int i=0;i<N;i++){
if(choose.count(groupChild[i])){
used.insert(i);
}
}
assert(used.size() != sum);
if(N > 2*used.size()){
now = 0;
fillResultSmall(group[0].second,group[0].first,0);
now = 0;
fillResult(group[1].second,group[1].first,u);
}else{
now = 0;
fillResultSmall(group[1].second,group[1].first,0);
now = 0;
fillResult(group[0].second,group[0].first,u);
}
}
}
}
vector<int> find_split(int n,int a,int b,int c,vector<int> p,vector<int> q){
N = n;
P = p;
Q = q;
group[0] = make_pair(a,1);
group[1] = make_pair(b,2);
group[2] = make_pair(c,3);
sort(group,group + 3);
for(int i=0;i<(int)p.size();i++){
adj[p[i]].push_back(q[i]);
adj[q[i]].push_back(p[i]);
}
vector<int> invalid(n);
res.resize(n);
dfsTree(0);
dfs(0,0);
dfsSmall(0,0);
if(!valid){
return invalid;
}
for(int i=0;i<n;i++){
if(!res[i]){
res[i] = group[2].second;
}
}
return res;
}