Submission #1099903

#	Time	Username	Problem	Language	Result	Execution time	Memory
1099903		model_code	Sphinx's Riddle (IOI24_sphinx)	C++17	100 / 100	343 ms	3272 KiB

sphinx

// correct/solution-author-refactor.cpp

#include <bits/stdc++.h>

#include "sphinx.h"
using namespace std;

const int MAX_N = 250 + 10;

const int ACTUAL_COLOR = -1;

using pii = pair<int, int>;

struct ComponentCounter {
  int n;
  vector<vector<int>> adjlist;

  int TT;
  vector<int> seen;

  ComponentCounter(vector<vector<int>> _adjlist)
      : n(_adjlist.size()), adjlist(_adjlist), TT(0), seen(n) {}
  ComponentCounter() : ComponentCounter(vector<vector<int>>()) {}

  /*
   * Marks all vertices in u's connected color component as seen, so long as
   * u's color is constant and known as per proposed_colors.
   *
   * proposed_colors is a list of colors that would be used to query.
   * u is a starting vertex
   */
  // Mark all as seen from u.
  void dfs(vector<int> &proposed_colors, int u) {
    seen[u] = TT;
    for (auto v : adjlist[u]) {
      if (proposed_colors[u] != ACTUAL_COLOR &&
          proposed_colors[u] == proposed_colors[v] && seen[v] != TT) {
        dfs(proposed_colors, v);
      }
    }
  }

  /*
   * Return number of connected color-components among vertices, according to
   * proposed_colors (only among known colors).
   */
  int num_components(vector<int> &proposed_colors, vector<int> &vertices) {
    TT++;
    int ret = 0;
    for (auto u : vertices) {
      if (seen[u] != TT) {
        ret++;
        dfs(proposed_colors, u);
      }
    }
    return ret;
  }
};

struct Groups {
  int n;
  vector<int> groupOf;
  vector<vector<int>> members;

  Groups() : Groups(0) {}
  Groups(int _n) : n(_n), groupOf(_n), members(_n) {
    for (int i = 0; i < n; i++) {
      groupOf[i] = i;
      members[i].push_back(i);
    }
  }

  int group_of(int x) { return groupOf[x]; }

  void join_groups_of(int x, int y) {
    int gx = group_of(x);
    int gy = group_of(y);

    if (members[gx].size() > members[gy].size()) {
      swap(gx, gy);
    }

    for (auto u : members[gx]) {
      groupOf[u] = gy;
      members[gy].push_back(u);
    }
    members[gx].clear();
  }

  vector<int> non_empty_groups() {
    vector<int> res;
    for (int i = 0; i < n; i++) {
      if (!members[i].empty()) {
        res.push_back(i);
      }
    }
    return res;
  }

  // Moves groups to the front
  void compress() {
    auto non_empty_group_indices = non_empty_groups();
    int upto = 0;
    for (auto g : non_empty_group_indices) {
      for (auto u : members[g]) {
        groupOf[u] = upto;
      }
      upto++;
    }
    for (int i = 0; i < n; i++) {
      members[i].clear();
    }
    for (int i = 0; i < n; i++) {
      members[groupOf[i]].push_back(i);
    }
  }

  /*
   * Facilitate a query on the "grouping" graph.
   *
   * This is equivalent to setting all vertices in a group to the group's color.
   * If group's color is ACTUAL_COLOR, then all vertices in the group will be
   * set to ACTUAL_COLOR. Note that this color must be the same for all vertices
   * in that group, since vertices in the same group belong to the same
   * color-connected component.
   *
   * Returns the number of color-connected components.
   * Note that this is always the same in the grouping graph as in the original
   * graph.
   */
  int queryGroups(vector<int> mask) {
    vector<int> proposed_colors(n);
    for (int i = 0; i < n; i++) {
      proposed_colors[i] = mask[groupOf[i]];
    }
    return perform_experiment(proposed_colors);
  }
};

/*
 * Identifies the color-connected components, but not their exact colors.
 */
struct Collapser {
  ComponentCounter componentCounter;
  int n;
  Groups grouper;

  Collapser(ComponentCounter _componentCounter)
      : componentCounter(_componentCounter),
        n(_componentCounter.n),
        grouper(n) {}

  // query for number of components in groups [s, e] union {cur}
  int queryOnly(vector<int> &groups, int s, int e, int cur) {
    vector<int> proposed_colors(n, n);
    for (int i = s; i < e; i++) {
      for (auto u : grouper.members[groups[i]]) {
        proposed_colors[u] = ACTUAL_COLOR;
      }
    }
    proposed_colors[cur] = ACTUAL_COLOR;

    vector<int> others;
    for (int i = 0; i < n; i++) {
      if (proposed_colors[i] == n) {
        others.push_back(i);
      }
    }
    int others_components =
        componentCounter.num_components(proposed_colors, others);

    auto found = perform_experiment(proposed_colors);
    auto res = found - others_components;
    return res;
  }

  // how many edges are from cur to vertices in groups [s, e) ?
  int extra_edges(vector<int> &groups, int s, int e, int cur) {
    auto if_none = e - s + 1;
    auto components = queryOnly(groups, s, e, cur);
    auto edges = if_none - components;
    return edges;
  }

  // there are num_edges edges from cur to vertices in groups [s,e).
  // find them and unionise with cur.
  void solve(vector<int> &groups, int s, int e, int num_edges, int cur) {
    if (num_edges == 0) return;
    if (e - s == num_edges) {
      for (int i = s; i < e; i++) {
        auto u = grouper.members[groups[i]].front();
        grouper.join_groups_of(cur, u);
      }
      return;
    }

    int mid = (e + s) / 2;

    auto left_edges = extra_edges(groups, s, mid, cur);
    solve(groups, s, mid, left_edges, cur);
    solve(groups, mid, e, num_edges - left_edges, cur);
  }

  void go() {
    for (int i = 1; i < n; i++) {
      auto non_empty_groups = grouper.non_empty_groups();
      vector<int> groups;
      for (auto g : non_empty_groups) {
        if (g < i) {
          groups.push_back(g);
        }
      }
      auto edges = extra_edges(groups, 0, groups.size(), i);
      solve(groups, 0, groups.size(), edges, i);
    }
  }

  Groups find_colors() {
    go();
    return grouper;
  }
};

struct TwoColor {
  vector<vector<int>> adjlist;
  vector<int> mark;

  // Defines a connected graph
  TwoColor(vector<vector<int>> _adjlist)
      : adjlist(_adjlist), mark(adjlist.size(), -1) {}

  void dfs(int u) {
    for (auto v : adjlist[u]) {
      if (mark[v] == -1) {
        mark[v] = mark[u] ^ 1;
        dfs(v);
      }
    }
  }

  // Returns a vector which is a 2-coloring of the graph, starting at 0.
  vector<int> two_color() {
    mark[0] = 0;
    dfs(0);
    return mark;
  }
};

// solution to subtask where graph is a coloring
namespace coloring {
Groups grouper;
ComponentCounter componentCounter;
vector<int> proposed_colors;

int known_colors[MAX_N];

// Find all of a particular color in subject, and write it to known_colors, and
// return it. Invariant: at least one of color in subject. Invariant:
// proposed_colors of all of test is color, proposed_colors of all of subject is
// ACTUAL_COLOR (or the actual color). here is number of components for
// proposed_colors
vector<int> find_of_color(int color, vector<int> &subject, vector<int> &test,
                          int here) {
  vector<int> ret;
  if (subject.size() <= 1) {
    for (auto u : subject) {
      known_colors[u] = color;
      ret.push_back(u);
    }
    return ret;
  }

  auto mid = subject.begin() + (subject.size() / 2);
  auto left = vector<int>(subject.begin(), mid);
  auto right = vector<int>(mid, subject.end());

  // check if any on left first
  for (auto u : right) {
    test.push_back(u);
    proposed_colors[u] = color;
  }
  int if_none =
      left.size() + componentCounter.num_components(proposed_colors, test);
  auto actual = grouper.queryGroups(proposed_colors);

  if (actual < if_none) {
    // there's some on the left
    ret = find_of_color(color, left, test, actual);
  }

  for (auto u : right) {
    test.pop_back();
    proposed_colors[u] = ACTUAL_COLOR;
  }

  // check if left accounts for everything
  for (auto u : left) {
    test.push_back(u);
    if (known_colors[u] == color) {
      proposed_colors[u] = color;
    }
  }
  auto test_components = componentCounter.num_components(proposed_colors, test);
  int if_none_right = right.size() + test_components;

  if (here < if_none_right) {
    // there's some on the right
    auto ret_right = find_of_color(color, right, test, here);
    for (auto u : ret_right) {
      ret.push_back(u);
    }
  }
  for (auto u : left) {
    proposed_colors[u] = ACTUAL_COLOR;
    test.pop_back();
  }

  return ret;
}

vector<int> find_colors(Groups _grouper, ComponentCounter _componentCounter,
                        int num_colors) {
  grouper = _grouper;
  componentCounter = _componentCounter;
  int n = componentCounter.n;
  for (int i = 0; i < n; i++) {
    proposed_colors.push_back(ACTUAL_COLOR);
    known_colors[i] = -1;
  }

  if (grouper.non_empty_groups().size() == 1) {
    for (int f = 0; f < num_colors; ++f) {
      vector<int> ord(num_colors, f);
      ord[0] = -1;
      if (perform_experiment(ord) == 1) return {f};
    }
  }

  vector<int> parts[2];
  int comp[2];

  TwoColor twoColor(componentCounter.adjlist);
  auto marks = twoColor.two_color();
  for (int i = 0; i < n; i++) {
    parts[marks[i]].push_back(i);
  }

  for (int z = 0; z < 2; z++) {
    for (auto u : parts[z]) {
      proposed_colors[u] = 0;
    }
    comp[z] = componentCounter.num_components(proposed_colors, parts[z]);
    for (auto u : parts[z]) {
      proposed_colors[u] = ACTUAL_COLOR;
    }
  }

  for (int c = 0; c < num_colors; c++) {
    for (int z = 0; z < 2; z++) {
      for (auto u : parts[z ^ 1]) {
        proposed_colors[u] = c;
      }
      int if_none = parts[z].size() + comp[z ^ 1];
      auto init = grouper.queryGroups(proposed_colors);

      if (init < if_none) {
        find_of_color(c, parts[z], parts[z ^ 1], init);
      }
      for (auto u : parts[z ^ 1]) {
        proposed_colors[u] = ACTUAL_COLOR;
      }
    }
  }

  return vector<int>(known_colors, known_colors + n);
}
}  // namespace coloring

vector<int> find_colours(int n, vector<int> a, vector<int> b) {
  vector<vector<int>> adjlist(n);
  int m = a.size();
  for (int i = 0; i < m; i++) {
    auto u = a[i];
    auto v = b[i];
    adjlist[u].push_back(v);
    adjlist[v].push_back(u);
  }
  ComponentCounter componentCounter(adjlist);

  // Returned solution has the property that every color component has a
  // distinct value. We use this for grouping.
  Collapser collapser(componentCounter);
  auto collapsed_groups = collapser.find_colors();
  collapsed_groups.compress();

  int collapsedNodes = collapsed_groups.non_empty_groups().size();
  vector<vector<int>> collapsedAdjlist(collapsedNodes);
  for (int u = 0; u < collapsedNodes; u++) {
    unordered_set<int> neighbors;
    for (auto uMember : collapsed_groups.members[u]) {
      for (auto vMember : adjlist[uMember]) {
        auto v = collapsed_groups.group_of(vMember);
        if (u != v) {
          neighbors.insert(v);
        }
      }
    }
    for (auto v : neighbors) {
      collapsedAdjlist[u].push_back(v);
    }
  }
  ComponentCounter collapsedCounter(collapsedAdjlist);

  auto groupedColors =
      coloring::find_colors(collapsed_groups, collapsedCounter, n);
  vector<int> known_colors;
  for (int i = 0; i < n; i++) {
    known_colors.push_back(groupedColors[collapsed_groups.group_of(i)]);
  }

  return known_colors;
}

• Subtask 0: 0 / 0
• Subtask 1: 3 / 3
• Subtask 2: 7 / 7
• Subtask 3: 33 / 33
• Subtask 4: 21 / 21
• Subtask 5: 36 / 36