#include "split.h" #include <bits/stdc++.h> using namespace std; #define pb push_back #define mp make_pair #define pii pair<int, int> #define fi first #define se second int n; bool found=0; vector<int> ans, sz, cc; vector<pii> vect; vector<vector<int> > graph; void dfs(int node, int p){ for (auto num:graph[node])if (num!=p)dfs(num, node), sz[node]+=sz[num]; } void label(int node, int p, int t){ if (cc[t]>=vect[t].fi)return; ++cc[t]; ans[node]=vect[t].se; for (auto num:graph[node])if (num!=p)label(num, node, t); } void dfs2(int node, int p){ if (found)return; for (auto num:graph[node])if (num!=p){ if (sz[num]>=vect[0].fi&&n-sz[num]>=vect[1].fi){ found=1; label(num, node, 0); label(node, num, 1); return; } else if (sz[num]>=vect[1].fi&&n-sz[num]>=vect[0].fi){ found=1; label(num, node, 1); label(node, num, 0); return; } else dfs2(num, node); } } vector<int> find_split(int N, int a, int b, int c, vector<int> p, vector<int> q){ n=N; cc.resize(4, 0); vect.pb(mp(a, 1)); vect.pb(mp(b, 2)); vect.pb(mp(c, 3)); sort(vect.begin(), vect.end()); ans.resize(n, 0); sz.resize(n, 1); graph.resize(n); for (int i=0; i<p.size(); ++i){ graph[p[i]].pb(q[i]); graph[q[i]].pb(p[i]); } dfs(0, -1); dfs2(0, -1); if (!found){ vector<int> temp(n, 0); return temp; } for (int i=0; i<n; ++i)if (!ans[i])ans[i]=vect[2].se; return ans; }