| # | 제출 시각 | 아이디 | 문제 | 언어 | 결과 | 실행 시간 | 메모리 |
|---|---|---|---|---|---|---|---|
| 1316220 | Kaguya | Split the Attractions (IOI19_split) | C++20 | 0 ms | 0 KiB |
vector<int> solve(int n, int a, int b, int c, vector<int> p, vector<int> q) {
int m = (int)p.size();
assert(n == m - 1);
vector<vector<int>> g(n);
for (int i = 0; i < n; i++) {
p[i]--; q[i]--;
g[p[i]].push_back(q[i]);
g[q[i]].push_back(p[i]);
}
vector<int> sz(n, 0);
vector<int> pr(n, -1);
function<void(int, int)> dfs = [&](int v, int pred) {
sz[v] = 1;
pr[v] = pred;
for (int to : g[v]) {
if (to == pred) {
continue;
}
dfs(to, v);
sz[v] += sz[to];
}
};
dfs(0, -1);
vector<pair<int, int>> ve = {{a, 0}, {b, 1}, {c, 2}};
sort(ve.begin(), ve.end());
vector<int> ans(n, -1);
auto color = [&](int from, int cnt, int col, int bad) -> void {
queue<int> que;
que.push(from);
vector<bool> was(n, false);
was[from] = true;
was[bad] = true;
que.push(from);
ans[from] = col;
int paint = 1;
while (!que.empty() && paint < cnt) {
int v = que.front();
que.pop();
for (int to : g[v]) {
if (was[to]) {
continue;
}
que.push(to);
was[to] = true;
if (paint < cnt) {
paint++;
ans[to] = col;
}
}
}
};
for (int i = 0; i < n; i++) {
if (pr[i] == -1) {
continue;
}
if (sz[i] >= ve[0].first && (n - sz[i]) >= ve[1].first) {
color(i, ve[0].first, ve[0].second, pr[i]);
color(pr[i], ve[1].first, ve[1].second, pr[i]);
}
}
for (int i = 0; i < n; i++) {
if (ans[i] == -1) {
ans[i] = ve.back().second;
}
}
return ans;
}