#include "split.h"
#include <bits/stdc++.h>
using namespace std;
template<typename T>
using vec = vector<T>;
using ll = long long;
using vll = vec<ll>;
using vvll = vec<vll>;
using pll = pair<ll,ll>;
#define pb push_back
#define dbg(x) cerr << #x << ": " << x << endl
vll subtree_size;
vec<int> color;
vec<pll> days;
vll dayOrder(3);
vvll adj;
pll good_edge = {-1,-1};
ll N;
void dfs(ll u, ll p) {
// cerr << "Call " << u << " " << p << endl;
subtree_size[u] = 1;
for(ll &v : adj[u]) {
if(v == p) continue;
dfs(v,u);
subtree_size[u] += subtree_size[v];
}
// can we cut (u,p) ?
if(subtree_size[u] >= days[0].first and N - subtree_size[u] >= days[1].first) {
// cerr << "Yay, " << u << " " << p << endl;
good_edge = {u,p};
}
else if(subtree_size[u] >= days[1].first and N - subtree_size[u] >= days[0].first) {
good_edge = {p,u};
}
}
void dfs_color(ll u, ll p, ll c) {
if(days[dayOrder[c-1]].first == 0) return;
color[u] = c;
days[dayOrder[c-1]].first--;
for(ll &v : adj[u]) {
if(v == p) continue;
dfs_color(v,u,c);
}
}
vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
N = n;
ll m = p.size();
adj.clear();
adj.resize(n);
for(int i = 0; i < m; i++) {
adj[p[i]].pb(q[i]);
adj[q[i]].pb(p[i]);
// cerr << "Edge " << p[i] << " " << q[i] << endl;
}
days = {{a,1}, {b,2}, {c,3}};
sort(days.begin(),days.end());
for(int i = 0; i < 3; i++) dayOrder[days[i].second - 1] = i;
subtree_size.assign(n,0);
color.assign(n,-1);
dfs(0,-1);
if(good_edge.first == -1) {
vec<int> out(n,0);
return out;
}
dfs_color(good_edge.first, good_edge.second, days[0].second);
// cerr << "Days: ";
// for(pll &x :days) cerr << "(" << x.first << ", " << x.second << "), ";
// cerr << endl;
// cerr << "Color so far: " << endl;
// for(int &x : color) cerr << x << " ";
// cerr << endl;
dfs_color(good_edge.second, good_edge.first, days[1].second);
for(int &x : color) x = (x != -1 ? x : days[2].second);
return color;
}
# | 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... |