#include "split.h" #include <bits/stdc++.h> using namespace std; #define pii pair<int, int> #define fr first #define se second const int MAXN = 1e5 + 10; const int INF = 0x3f3f3f3f; vector<int> adj[MAXN]; set<int> child[MAXN]; vector<int> ans; int dp[MAXN]; int l, r; bool in(int x) { return ((l <= x) && (x <= r)); } int ok = -1; void calc(int node, int dad) { dp[node] = 1; for (auto u : adj[node]) { if (u == dad) continue; calc(u, node); dp[node] += dp[u]; child[node].insert(u); } if (in(dp[node])) ok = node; } bool add(int a, int b) { // a -> b dp[a] += dp[b]; child[a].insert(b); return in(dp[a]);; } bool re(int a, int b) { // remove a -> b dp[a] -= dp[b]; child[a].erase(b); return in(dp[a]);; } bool reroot(int node, int dad) { for (auto u : adj[node]) { if (u == dad) continue; bool x = re(node, u); add(u, node); if (x) { ok = node; return true; } if (reroot(u, node)) return true; x = re(u, node); add(node, u); if (x) { ok = u; return true; } } return false; } int pr[MAXN]; int qChild[MAXN]; vector<int> subT; void findSub(int node, int dad) { subT.push_back(node); for (auto u : adj[node]) { if (u == dad) continue; findSub(u, node); } } int qtd, qual; void fillg(int node, int dad) { ans[node] = qual; qtd--; for (auto u : adj[node]) if ((u != dad) && (qtd > 0)) fillg(u, node); } vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) { vector<pii> ord = {{a, 1}, {b, 2}, {c, 3}}; sort(ord.begin(), ord.end()); l = ord[0].fr; r = (n - ord[1].fr); for (int i = 0; i < ((int) p.size()); i++) { adj[p[i]].push_back(q[i]); adj[q[i]].push_back(p[i]); } ans.resize(n); // calc calc(0, -1); if (ok == -1) reroot(0, -1); if (ok == -1) return ans; // return empty ans // we are going to remove the extra ones in the min, and it's going to work out, by going to the leaves for (int i = 0; i < n; i++) { for (auto u : child[i]) pr[u] = i; qChild[i] = child[i].size(); } findSub(ok, pr[ok]); vector<int> zeroes; for (auto u : subT) { if (qChild[u] == 0) zeroes.push_back(u); ans[u] = ord[0].se; } int qover = (dp[ok] - ord[0].fr); while ((qover > 0) && !zeroes.empty()) { int u = zeroes.back(); zeroes.pop_back(); qover--; ans[u] = ord[2].se; qChild[pr[u]]--; if (qChild[pr[u]] == 0) zeroes.push_back(pr[u]); } qtd = ord[1].fr; qual = ord[1].se; fillg(pr[ok], ok); // retornar for (auto& u : ans) if (u == 0) u = ord[2].se; return ans; }