Submission #223701

#	Time	Username	Problem	Language	Result	Execution time	Memory
223701		rama_pang	Split the Attractions (IOI19_split)	C++14	100 / 100	150 ms	18024 KiB

split

#include "split.h"
#include <bits/stdc++.h>
using namespace std;

class DisjointSet {
 private:
  vector<int> p, sz;

 public:
  DisjointSet(int n) {
    p.assign(n, 0);
    iota(begin(p), end(p), 0);
    sz.assign(n, 1);
  }

  int Find(int x) {
    return p[x] == x ? x : p[x] = Find(p[x]);
  }

  bool Unite(int x, int y) {
    x = Find(x), y = Find(y);
    if (x == y) return false;
    if (sz[x] > sz[y]) swap(x, y);
    p[x] = y;
    sz[y] += sz[x];
    return true;
  }

  int Size(int x) {
    return sz[Find(x)];
  }

  bool SameSet(int x, int y) {
    return Find(x) == Find(y);
  }
};

int find_root(int mxsize, const vector<vector<int>> &adj) {
  vector<int> sz(adj.size());
  vector<int> par(adj.size(), -1);

  function<void(int)> dfs = [&](int u) {
    sz[u] = 1;
    for (auto v : adj[u]) if (v != par[u]) {
      par[v] = u;
      dfs(v);
      sz[u] += sz[v];
    }
  };

  dfs(0);

  int u = 0;
  while (true) {
    pair<int, int> mx = {-1, -1};
    for (auto v : adj[u]) if (v != par[u]) {
      mx = max(mx, {sz[v], v});
    }
    if (mx.first <= mxsize) return u;
    u = mx.second;
  }
}

vector<int> find_split(int n, int a, int b, int c, 
                       vector<int> p, vector<int> q) {
  DisjointSet d(n);
  vector<int> splits(n, 0);
  int m = p.size();
  int ca, cb, cc; // color of a, b, c

  { // recolor a, b, c such that a <= b <= c holds
    vector<pair<int, int>> split_sizes = {{a, 1}, {b, 2}, {c, 3}};
    sort(begin(split_sizes), end(split_sizes));
    tie(a, ca) = split_sizes[0];
    tie(b, cb) = split_sizes[1];
    tie(c, cc) = split_sizes[2];
  }

  vector<vector<int>> adj(n); // spanning tree adjacency list
  for (int i = 0; i < m; i++) {
    if (d.Unite(p[i], q[i])) {
      adj[p[i]].emplace_back(q[i]);
      adj[q[i]].emplace_back(p[i]);
    }
  }

  d = DisjointSet(n);
  int r = find_root(n - b, adj);

  // Unite all subtrees of children of root
  for (int i = 0; i < n; i++) {
    for (auto j : adj[i]) {
      if (i == r || j == r) continue;
      d.Unite(i, j);
    }
  }

  for (int i = 0; i < m; i++) {
    if (d.SameSet(p[i], q[i])) continue;
    if (p[i] == r || q[i] == r) continue;
    if (d.Size(p[i]) >= a || d.Size(q[i]) >= a) continue;
    d.Unite(p[i], q[i]);
    adj[p[i]].emplace_back(q[i]);
    adj[q[i]].emplace_back(p[i]);
  }
  
  function<void(int, int&, int, int)> color = [&](int u, int &n, int c, int bad) {
    if (u == bad || splits[u] != 0 || n == 0) return;
    splits[u] = c, n--;
    for (auto v : adj[u]) color(v, n, c, bad);
  };

  for (auto v : adj[r]) if (d.Size(v) >= a) {
    color(v, a, ca, r); // color merged subtree for set A
    color(r, b, cb, -1); // color root and its surrounding subgraph for set B
    replace(begin(splits), end(splits), 0, cc); // color vertices not in A or B for set C
    break;
  }
  
  return splits;
}

• Subtask 1: 7 / 7
• Subtask 2: 11 / 11
• Subtask 3: 22 / 22
• Subtask 4: 24 / 24
• Subtask 5: 36 / 36