#include<iostream>
#include<vector>
#include<algorithm>
#include<queue>
#include "split.h"
using namespace std;
const int MAX_N=2e5+5;
vector<int>g[MAX_N];
int col[MAX_N];
int deg[MAX_N];
int A,B;
int whA,whB;
int parent[MAX_N];
int sz[MAX_N];
void dfs(int u)
{
sz[u]=1;
for(int v:g[u])
{
if(v==parent[u])continue;
deg[u]++;
deg[v]++;
parent[v]=u;
dfs(v);
sz[u]+=sz[v];
}
}
void colour(int u,int color)
{
col[u]=color;
for(int v:g[u])
{
if(v==parent[u])continue;
colour(v,color);
}
}
vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q)
{
vector<int>ans;
ans.resize(n);
vector<pair<int,int>>shit;
shit.push_back({a,1});shit.push_back({b,2});shit.push_back({c,3});
sort(shit.begin(),shit.end());
A=shit[0].first;whA=shit[0].second;
B=shit[1].first;whB=shit[1].second;
for(int i=0;i<p.size();i++)
{
g[p[i]].push_back(q[i]);
g[q[i]].push_back(p[i]);
}
parent[0]=-1;
dfs(0);
vector<pair<int,int>>subtrees;
for(int u:g[0])
{
subtrees.push_back({sz[u],u});
}
sort(subtrees.begin(),subtrees.end());
int firstA=-1,firstB=-1;
for(auto [s,u]:subtrees)
{
if(s>=A && firstA==-1)
{
firstA=u;
}
if(s>=B && firstB==-1)
{
firstB=u;
}
}
if(firstA==-1)return ans;
if(n-sz[firstA]>=B)
{
colour(0,whB);
colour(firstA,whA);
}
else if(firstB!=-1 && n-sz[firstB]>=A)
{
colour(0,whA);
colour(firstB,whB);
}
else return ans;
queue<int>bfs;
int cntwhA=0;
for(int i=0;i<n;i++)
{
if(col[i]==whA)
{
cntwhA++;
if(deg[i]==1)bfs.push(i);
}
}
int toremove=cntwhA-A;
while(bfs.size() && toremove>0)
{
int u=bfs.front();
bfs.pop();
toremove--;
col[u]=shit[2].second;
for(int v:g[u])
{
if(col[v]==whA)
{
deg[v]--;
if(deg[v]==1)bfs.push(v);
}
}
}
while(bfs.size())bfs.pop();
int cntwhB=0;
for(int i=0;i<n;i++)
{
if(col[i]==whB)
{
cntwhB++;
if(deg[i]==1)bfs.push(i);
}
}
toremove=cntwhB-B;
while(bfs.size() && toremove>0)
{
int u=bfs.front();
bfs.pop();
toremove--;
col[u]=shit[2].second;
for(int v:g[u])
{
if(col[v]==whB)
{
deg[v]--;
if(deg[v]==1)bfs.push(v);
}
}
}
//B
for(int i=0;i<n;i++)ans[i]=col[i];
return ans;
}
| # | 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... |