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>
using namespace std;
typedef long long ll;
int n, a, b, c, sz[100100], par[100100], idx[5];
void dfs(int s, int &cnt, int no, const vector<vector<int>> &adj, int num, vector<int> &ret){
if (!cnt) return;
--cnt;
ret[s] = num;
if (no==-1) sz[s] = 1;
for (auto &v:adj[s]) if (v!=no && ret[v]!=num){
if (no==-1) par[v] = s;
dfs(v, cnt, no, adj, num, ret);
if (no==-1) sz[s] += sz[v];
}
}
vector<int> solve_tree(const vector<vector<int>> &adj){
//printf("n = %d\n", n);
vector<int> ret(n);
int N = n;
dfs(0, N, -1, adj, idx[3], ret);
for (int i=1;i<n;i++){
//printf("%d: %d\n", i, sz[i]);
int x = sz[i], y = n - sz[i];
if (a<=x && b<=y){
dfs(i, a, par[i], adj, idx[1], ret);
dfs(par[i], b, i, adj, idx[2], ret);
return ret;
}
if (a<=y && b<=x){
dfs(par[i], a, i, adj, idx[1], ret);
dfs(i, b, par[i], adj, idx[2], ret);
return ret;
}
}
fill(ret.begin(), ret.end(), 0);
return ret;
}
vector<int> find_split(int N, int A, int B, int C, vector<int> p, vector<int> q) {
n = N;
vector<pair<int, int>> tmp;
tmp.emplace_back(A, 1);
tmp.emplace_back(B, 2);
tmp.emplace_back(C, 3);
sort(tmp.begin(), tmp.end());
a = tmp[0].first, b = tmp[1].first, c = tmp[2].first;
idx[1] = tmp[0].second, idx[2] = tmp[1].second, idx[3] = tmp[2].second;
vector<vector<int>> adj0(N, vector<int>());
for (int i=0;i<(int)p.size();i++){
adj0[p[i]].push_back(q[i]);
adj0[q[i]].push_back(p[i]);
}
auto ans = solve_tree(adj0);
int cnt[4] = {0};
for (auto &x:ans) cnt[x]++;
assert(cnt[1]==A);
assert(cnt[2]==B);
assert(cnt[3]==C);
return ans;
}
# | 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... |