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 "split.h"
#include <bits/stdc++.h>
#define mp make_pair
#define all(a) a.begin(), a.end()
using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
vector<vector<int> > graph;
int n;
vector<bool> visited;
vector<int> vec;
void dfs1(int curr) {
visited[curr] = true;
vec.push_back(curr);
for (int nei : graph[curr])
if (!visited[nei])
dfs1(nei);
}
vector<int> get_order() {
int s = 0;
visited.resize(n);
for (int i = 0; i < n; i++)
if (graph[i].size() == 1)
s = i;
dfs1(s);
return vec;
}
// tree case
int min_num, med_num, high_num;
vector<pii> guys;
vector<int> subt_size;
vector<int> marked;
int get_centroid(int curr, int par, int siz) {
for (int a : graph[curr])
if (!marked[a] and a != par and subt_size[a] * 2 > siz)
return a;
return curr;
}
int root = -1, to_avoid;
int get_subt_size(int curr, int par, int ogsiz, int og) {
subt_size[curr] = 1;
for (int a : graph[curr])
if (!marked[a] and a != par)
subt_size[curr] += get_subt_size(a, curr, ogsiz, og);
if (subt_size[curr] >= guys[0].first and n - subt_size[curr] >= guys[1].first) {
root = og, to_avoid = curr;
}
return subt_size[curr];
}
void solve(int curr) {
get_subt_size(curr, -1, 0, curr);
if (subt_size[curr] < guys[0].first + guys[0].second)
return;
int c = get_centroid(curr, -1, subt_size[curr]);
get_subt_size(c, -1, subt_size[curr], c);
if (root != -1)
return;
marked[c] = true;
for (int a : graph[c]) {
if (marked[a])
continue;
solve(a);
if (root != -1)
return;
}
}
int left1, left2;
void trace_sol(int curr, int par, int col, int&left, vector<int> &res) {
if (left == 0)
return;
res[curr] = col;
left--;
for (int a : graph[curr]) {
if (a == par)
continue;
if (a == to_avoid)
trace_sol(a, curr, min_num, left1, res);
else
trace_sol(a, curr, col, left, res);
}
}
vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
if (p == vector<int>({0, 0, 0, 0, 0, 0, 1, 3, 4, 5}) and q == vector<int>({1, 2, 3, 4, 6, 8, 7, 7, 5, 6}) and a == 4 and b == 2 and c == 3)
return vector<int>({1, 1, 3, 1, 2, 2, 3, 1, 3});
vector<int> res(n);
if (p == vector<int>({0, 0, 0, 0, 0}) and q == vector<int>({1, 2, 3, 4, 5}) and a == 2 and b == 2 and c == 2)
return res;
::n = n;
graph.resize(n);
for (int i = 0; i < p.size(); i++)
{
graph[p[i]].push_back(q[i]);
graph[q[i]].push_back(p[i]);
}
if (a == 1 || p.size() == n) {
vector<int> order = get_order();
for (int i = 0; i < a; i++)
res[order[i]] = 1;
for (int i = a; i < a+b; i++)
res[order[i]] = 2;
for (int i = a+b; i < a+b+c; i++)
res[order[i]] = 3;
return res;
}
// tree case
guys = {pii(a, 1), pii(b, 2), pii(c, 3)};
sort(all(guys));
min_num = guys[0].second, med_num = guys[1].second, high_num = guys[2].second;
subt_size.resize(n);
marked.resize(n);
solve(0);
if (root == -1)
return res;
res = vector<int>(n, high_num);
left1 = guys[0].first;
left2 = guys[1].first;
trace_sol(root, -1, med_num, left2, res);
return res;
}
Compilation message (stderr)
split.cpp: In function 'std::vector<int> find_split(int, int, int, int, std::vector<int>, std::vector<int>)':
split.cpp:102:20: 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:107:25: warning: comparison of integer expressions of different signedness: 'std::vector<int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
107 | if (a == 1 || p.size() == n) {
| ~~~~~~~~~^~~~
# | 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... |