Submission #1279562

#TimeUsernameProblemLanguageResultExecution timeMemory
1279562avighnaSphinx's Riddle (IOI24_sphinx)C++20
43 / 100
342 ms1596 KiB
#include <bits/stdc++.h>

int perform_experiment(std::vector<int> E);

int n;
std::vector<std::vector<int>> adj;

int query(const std::vector<int> &c) {
  if (std::find(c.begin(), c.end(), -1) != c.end()) {
    return perform_experiment(c);
  }
  int ans = 0;
  std::queue<int> q;
  std::vector<bool> vis(n);
  for (int i = 0; i < n; ++i) {
    if (vis[i]) {
      continue;
    }
    vis[i] = true;
    ans++;
    q.push(i);
    while (!q.empty()) {
      int u = q.front();
      q.pop();
      for (int &i : adj[u]) {
        if (!vis[i] and c[u] == c[i]) {
          vis[i] = true;
          q.push(i);
        }
      }
    }
  }
  return ans;
}

class dsu {
private:
  int n;
  std::vector<int> par, size;
  std::vector<std::vector<int>> comp_m;

public:
  dsu(int n) : n(n), par(n), size(n, 1), comp_m(n) {
    std::iota(par.begin(), par.end(), 0);
    for (int i = 0; i < n; ++i) {
      comp_m[i].push_back(i);
    }
  }

  int root(int u) { return u == par[u] ? u : par[u] = root(par[u]); }

  void merge(int u, int v) {
    u = root(u), v = root(v);
    if (u == v) {
      return;
    }
    if (size[u] < size[v]) {
      std::swap(u, v);
    }
    for (int &i : comp_m[v]) {
      comp_m[u].push_back(i);
    }
    par[v] = u;
    size[u] += size[v];
  }

  const std::vector<int> &comp(int u) { return comp_m[root(u)]; }
};

std::vector<int> find_colours(int N, std::vector<int> x, std::vector<int> y) {
  n = N;
  adj.resize(n);
  for (int i = 0; i < x.size(); ++i) {
    adj[x[i]].push_back(y[i]);
    adj[y[i]].push_back(x[i]);
  }

  dsu dsu(n);
  for (int u = 1; u < n; ++u) {
    std::vector<bool> seen(n);
    std::vector<int> nodes;
    for (int i : adj[u]) {
      i = dsu.root(i);
      if (i > u or seen[i]) {
        continue;
      }
      seen[i] = true;
      nodes.push_back(i);
    }
    auto has = [&](int from, int till) {
      std::vector<int> q(n, n);
      for (int i = from; i <= till; ++i) {
        for (int u : dsu.comp(nodes[i])) {
          q[u] = dsu.root(nodes[i]);
        }
      }
      q[u] = n + 1;
      int bef = query(q);
      for (int i = from; i <= till; ++i) {
        for (int u : dsu.comp(nodes[i])) {
          q[u] = -1;
        }
      }
      q[u] = -1;
      int aft = query(q);
      return bef - aft;
    };
    int times = has(0, nodes.size() - 1);
    for (int _ = 0, fr = -1; _ < times; ++_) {
      fr = *std::ranges::partition_point(std::views::iota(fr + 1, int(nodes.size())), [&](int i) {
        return has(fr + 1, i) == 0;
      });
      dsu.merge(u, nodes[fr]);
    }
  }

  std::vector<std::vector<int>> nadj(n);
  for (int u = 0; u < n; ++u) {
    for (int &i : adj[u]) {
      if (dsu.root(u) != dsu.root(i)) {
        nadj[dsu.root(u)].push_back(dsu.root(i));
      }
    }
  }
  for (int u = 0; u < n; ++u) {
    std::sort(nadj[u].begin(), nadj[u].end());
    nadj[u].erase(std::unique(nadj[u].begin(), nadj[u].end()), nadj[u].end());
  }

  std::vector<int> dist(n);
  std::queue<int> q;
  std::vector<bool> vis(n);
  for (int i = 0; i < n; ++i) {
    if (vis[i] or dsu.root(i) != i) {
      continue;
    }
    vis[i] = true;
    q.push(i);
    while (!q.empty()) {
      int u = q.front();
      q.pop();
      for (int &i : nadj[u]) {
        if (!vis[i]) {
          vis[i] = true;
          dist[i] = dist[u] + 1;
          q.push(i);
        }
      }
    }
  }

  std::vector<int> a, b;
  for (int u = 0; u < n; ++u) {
    if (dsu.root(u) != u) {
      continue;
    }
    (dist[u] % 2 ? a : b).push_back(u);
  }

  std::vector<int> ans(n);
  assert(!b.empty());
  if (a.empty()) {
    for (int c = 0; c < n; ++c) {
      std::vector<int> q(n, -1);
      q[0] = c;
      if (query(q) == 1) {
        ans[b[0]] = c;
        break;
      }
    }
  }

  if (!a.empty()) {
    for (auto &[a, b] : std::array<std::pair<std::vector<int>, std::vector<int>>, 2>({{a, b}, {b, a}})) {
      for (int c = 0; c < n; ++c) {
        auto has = [&](int from, int till) {
          std::vector<int> q(n, n);
          for (int &i : b) {
            for (auto &u : dsu.comp(i)) {
              q[u] = n + 1;
            }
          }
          for (int i = from; i <= till; ++i) {
            for (auto &u : dsu.comp(a[i])) {
              q[u] = a[i];
            }
          }
          int bef = query(q);
          for (int &i : b) {
            for (auto &u : dsu.comp(i)) {
              q[u] = c;
            }
          }
          for (int i = from; i <= till; ++i) {
            for (auto &u : dsu.comp(a[i])) {
              q[u] = -1;
            }
          }
          int aft = query(q);
          return bef - aft;
        };
        int fr = -1;
        while (true) {
          if (has(fr + 1, a.size() - 1) == 0) {
            break;
          }
          fr = *std::ranges::partition_point(std::views::iota(fr + 1, int(a.size())), [&](int i) {
            return has(fr + 1, i) == 0;
          });
          if (fr == a.size()) {
            break;
          }
          ans[a[fr]] = c;
        }
      }
    }
  }

  for (int i = 0; i < n; ++i) {
    ans[i] = ans[dsu.root(i)];
  }
  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...
#Verdict Execution timeMemoryGrader output
Fetching results...