Submission #1316231

#TimeUsernameProblemLanguageResultExecution timeMemory
1316231KaguyaSplit the Attractions (IOI19_split)C++20
0 / 100
1 ms336 KiB
#include <bits/stdc++.h>
using namespace std;

vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
  int m = (int)p.size();
  assert(m == n - 1);
  vector<vector<int>> g(n);
  for (int i = 0; i < m; 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, 1}, {b, 2}, {c, 3}};
  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, i);
    }
  }
  for (int i = 0; i < n; i++) {
    if (ans[i] == -1) {
      ans[i] = ve.back().second;
    }
  }
  return ans;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...