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>
#define INF 10000000000000000
#define MOD 1000000017
#define mid(l, u) ((l+u)/2)
#define rchild(i) (i*2 + 2)
#define lchild(i) (i*2 + 1)
#define rr mt_rand()
using namespace std;
int sz[100005];
bool hasBackEdge[100005];
vector<pair<int, int>> t;
int piv = -1;
bool ansfound = false;
vector<int> sgive;
int ss, rs;
vector<int> o;
pair<int, int> se[100005];
bool keep[100005];
vector<int> tadj[100005];
/*
9 4 2 3 10
0 1
0 2
0 3
0 4
0 6
0 8
1 7
3 7
4 5
5 6
*/
int dfs(int i, vector<int> adj[], int vis[]){
if(vis[i]==1 || vis[i]==2) return 0;
se[i].first = o.size();
o.push_back(i);
vis[i] = 1;
sz[i] = 1;
bool found = false;
for(int j: adj[i]){
if(vis[j]==0){
tadj[i].push_back(j);
//cout<<i<<" "<<j<<endl;
}
//cout<<"calling from "<<i<<endl;
sz[i]+=dfs(j, adj, vis);
if(sz[j]>=t[0].first) found = true;
}
if(sz[i]>=t[0].first && !found){
//cout<<"pivot at "<<i<<endl;
piv = i;
}
vis[i] = 2;
se[i].second = o.size();
o.push_back(i);
return sz[i];
}
bool dfs1(int i, vector<int> adj[], int vis[]){
//cout<<"visiting "<<i<<endl;
if(se[i].first < se[piv].first || se[i].first>se[piv].second){
//cout<<i<<" is before the pivot"<<endl;
return true;
}
if(vis[i]==1 || vis[i]==2) return false;
vis[i] = 1;
hasBackEdge[i] = false;
for(int j: adj[i]){
//cout<<"trying to visit "<<j<<endl;
hasBackEdge[i] = dfs1(j, adj, vis) || hasBackEdge[i];
}
vis[i] = true;
return hasBackEdge[i];
}
void dfs2(int i){
//cout<<"Visiting "<<i<<endl;
if((ss - sz[i]) >= t[0].first && hasBackEdge[i]){
//cout<<"Removing "<<i<<endl;
keep[i] = false;
ss-=sz[i];
rs+=sz[i];
sgive.push_back(i);
return;
}
for(int j: tadj[i]){
dfs2(j);
}
}
int rem;
void dfs3(int i, int lbl, vector<int> &ans){ //assign from the remaining pile
if(rem==0) return;
if(i==piv) return;
ans[i] = lbl;
rem--;
for(int j: tadj[i]){
dfs3(j, lbl, ans);
}
}
void dfs4(int i, int lbl, vector<int> &ans){
if(rem==0) return;
if(!keep[i]) return;
ans[i] = lbl;
rem--;
for(int j: tadj[i]){
dfs4(j, lbl, ans);
}
}
vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q){
map<int, int> mp;
t = {{a, 1}, {b, 2}, {c, 3}};
sort(t.begin(), t.end());
vector<int> adj[n];
for(int i = 0;i<p.size();i++){
adj[p[i]].push_back(q[i]);
adj[q[i]].push_back(p[i]);
}
int vis[n];
for(int i =0 ;i<n;i++){
vis[i] = 0;
}
dfs(0, adj, vis);
for(int i =0 ;i<n;i++){
vis[i] = 0;
hasBackEdge[i] = false;
}
dfs1(piv, adj, vis);
ss = sz[piv];
rs = n - sz[piv];
vector<int> ans;
for(int i = 0;i<n;i++){
ans.push_back(0);
keep[i] = true;
}
dfs2(piv);
if(rs<t[0].first || ss<t[0].first) return ans;
int lblr, lbls;
int szr, szs;
sgive.push_back(0);
if(rs>=t[1].first){
lblr = t[1].second;
szr = t[1].first;
lbls = t[0].second;
szs = t[0].first;
}
else{
lblr = t[0].second;
szr = t[0].first;
lbls = t[1].second;
szs = t[1].first;
}
rem = szr;
for(int j: sgive){
dfs3(j, lblr, ans);
}
rem = szs;
dfs4(piv, lbls, ans);
for(int i = 0;i<n;i++){
if(ans[i]==0) ans[i] = t[2].second;
}
return ans;
}
/*signed main(){
int n, a, b, c, m;
cin>>n>>a>>b>>c>>m;
vector<int> p, q;
for(int i = 0;i<m;i++){
int u, v;
cin>>u>>v;
p.push_back(u);
q.push_back(v);
}
vector<int> ans = find_split(n, a, b, c, p, q);
for(int j: ans) cout<<j<<" ";
cout<<endl;
}*/
//l, mid(l, r) || l+1, r
//
Compilation message (stderr)
split.cpp: In function 'std::vector<int> find_split(int, int, int, int, std::vector<int>, std::vector<int>)':
split.cpp:112:20: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
for(int i = 0;i<p.size();i++){
~^~~~~~~~~
# | 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... |