#include "split.h"
#include <bits/stdc++.h>
using namespace std;
const int nx=2e5+5;
int sz[nx], c, h[nx], szc[nx], cnt, sm, used[nx], vs[nx], fans;
vector<pair<int, int>> d[nx], t;
vector<int> lst, g[nx], res;
pair<int, int> mx;
int dfssz(int u)
{
sz[u]=1;
for (auto [v, idx]:d[u]) if (!sz[v]) used[idx]=1, sz[u]+=dfssz(v);
return sz[u];
}
void dfsfill(int u, int p, int hd)
{
h[u]=hd;
szc[hd]++;
for (auto [v, idx]:d[u]) if (v!=p&&used[idx]) dfsfill(v, u, hd);
}
void fillans(int u, int p, int ans)
{
if (cnt==0) return;
res[u]=ans;
cnt--;
for (auto [v, idx]:d[u]) if (v!=p&&v!=c&&!res[v]&&used[idx]) fillans(v, u, ans);
}
int findcentroid(int u, int p, int rtsz)
{
for (auto [v, idx]:d[u]) if (v!=p&&2*sz[v]>rtsz&&used[idx]) return findcentroid(v, u, rtsz);
return u;
}
void dfsfind(int u, int nd)
{
if (fans) return;
if (sm+szc[u]>=t[0].first)
{
fans=1;
for (auto v:lst) cnt=szc[v], fillans(v, c, t[0].second);
cnt=t[0].first-sm;
fillans(nd, nd, t[0].second);
cnt=t[1].first;
fillans(c, c, t[1].second);
for (auto &x:res) if (x==0) x=t[2].second;
return;
}
vs[u]=1;
sm+=szc[u];
lst.push_back(u);
for (auto v:g[u]) if (!vs[h[v]]) dfsfind(h[v], v);
}
vector<int> find_split(int n, int A, int B, int C, vector<int> p, vector<int> q) {
res.resize(n);
t.push_back({A, 1});
t.push_back({B, 2});
t.push_back({C, 3});
sort(t.begin(), t.end());
for (int i=0; i<p.size(); i++) d[p[i]].push_back({q[i], i}), d[q[i]].push_back({p[i], i});
dfssz(0);
c=findcentroid(0, 0, n);
for (auto [v, idx]:d[c])
{
if (!used[idx]) continue;
dfsfill(v, c, v);
mx=max(mx, {szc[v], v});
}
for (int i=0; i<p.size(); i++) if (!used[i]&&p[i]!=c&&q[i]!=c&&h[p[i]]!=h[q[i]]) g[h[p[i]]].push_back(q[i]), g[h[q[i]]].push_back(p[i]);
if (mx.first>=t[0].first)
{
cnt=t[0].first;
fillans(mx.second, c, t[0].second);
cnt=t[1].first;
fillans(c, mx.second, t[1].second);
for (int i=0; i<n;i ++) if (!res[i]) res[i]=t[2].second;
return res;
}
for (auto [v, idx]:d[c])
{
if (vs[v]||!used[idx]) continue;
lst.clear();
sm=0;
dfsfind(v, v);
if (fans) return res;
}
return res;
}
# | 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... |