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>
using namespace std;
const int N = 1e5+5;
vector<int> tree[N];
int subtreesize[N]={};
int par[N]={};
vector<int> ans;
const int max_afford = 3e6;
int go(int v, int p=0) {
par[v]=p;
for (int i = 0; i < tree[v].size(); ++i) {
int u = tree[v][i];
if (u != p) subtreesize[v] += go(u, v);
}
return subtreesize[v];
}
int cursz;
void fill_ans(int v, int p, int num) {
if (ans[v] != 0) return;
if (cursz == 0) return;
cursz--;
ans[v] = num;
for (int i = 0; i < tree[v].size(); ++i) {
int u = tree[v][i];
if (u != p) fill_ans(u, v, num);
}
}
bool solve_for_tree(const vector<pair<int, int>>& edges, int sz1, int sz2, int num1, int num2, int root = 0) {
const int n = edges.size()+1;
for (int i = 0; i < n; ++i) tree[i].clear(), subtreesize[i]=1;
for (int i = 0; i < n-1; ++i) {
tree[edges[i].first].push_back(edges[i].second);
tree[edges[i].second].push_back(edges[i].first);
}
go(root, root);
for (int i = 0; i < n; ++i) {
int here = subtreesize[i];
int rest = n-subtreesize[i];
if (here >= sz1 && rest >= sz2) {
cursz = sz1;
fill_ans(i, par[i], num1);
cursz = sz2;
fill_ans(root, root, num2);
for (int j = 0; j < n; ++j) {
if (ans[j] == 0) ans[j] = 6 - num1 - num2;
}
return true;
}
if (here >= sz2 && rest >= sz1) {
cursz = sz2;
fill_ans(i, par[i], num2);
cursz = sz1;
fill_ans(root, root, num1);
for (int j = 0; j < n; ++j) {
if (ans[j] == 0) ans[j] = 6 - num1 - num2;
}
return true;
}
}
return false;
}
vector<pair<int, int>> e;
vector<pair<int, int>> all_e;
struct DSU {
vector<int> p;
vector<int> sz;
void init(int n) {
p.resize(n);
sz.resize(n);
p.assign(n, 1);
sz.assign(n, 1);
iota(p.begin(), p.end(), 0);
}
DSU(int n) {
init(n);
}
int get(int a) {
return p[a] = (a == p[a] ? a : get(p[a]));
}
void merge(int a, int b) {
int x = get(a), y = get(b);
if (x == y) return;
if (sz[x] > sz[y]) swap(x, y);
p[x] = y;
sz[y] += sz[x];
}
};
vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
ans.resize(n, 0);
int its = max_afford / n;
if (n <= 2.5e3) {
its /= n;
its *= 3;
}
for (int i = 0; i < p.size(); ++i) {
all_e.push_back({p[i], q[i]});
}
vector<pair<int, int>> sizes={{a,1}, {b,2}, {c, 3}};
sort(sizes.begin(), sizes.end());
DSU d(0);
for (int _ = 0; _ < its; ++_) {
d.init(n);
e.clear();
e.shrink_to_fit();
random_shuffle(all_e.begin(), all_e.end());
int cnt = 0;
for (int i = 0; i < all_e.size(); ++i) {
if (d.get(all_e[i].first) != d.get(all_e[i].second)) {
e.push_back(all_e[i]);
d.merge(all_e[i].first, all_e[i].second);
cnt++;
if (cnt == n-1) break;
}
}
bool ok = false;
if (n > 2.5e3)
ok |= solve_for_tree(e, sizes[0].first, sizes[1].first, sizes[0].second, sizes[1].second);
else {
for (int i = 0; i < n; ++i) {
ok |= solve_for_tree(e, sizes[0].first, sizes[1].first, sizes[0].second, sizes[1].second, i);
}
}
if (ok) break;
}
return ans;
}
Compilation message (stderr)
split.cpp: In function 'int go(int, int)':
split.cpp:14:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
14 | for (int i = 0; i < tree[v].size(); ++i) {
| ~~^~~~~~~~~~~~~~~~
split.cpp: In function 'void fill_ans(int, int, int)':
split.cpp:26:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
26 | for (int i = 0; i < tree[v].size(); ++i) {
| ~~^~~~~~~~~~~~~~~~
split.cpp: In function 'std::vector<int> find_split(int, int, int, int, std::vector<int>, std::vector<int>)':
split.cpp:102:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
102 | for (int i = 0; i < p.size(); ++i) {
| ~~^~~~~~~~~~
split.cpp:114:27: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
114 | for (int i = 0; i < all_e.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... |