Submission #1040177

#	Time	Username	Problem	Language	Result	Execution time	Memory
1040177		spacewalker	Keys (IOI21_keys)	C++17	100 / 100	289 ms	210568 KiB

keys

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

template<class T, class U>
ostream& operator<< (ostream& os, const pair<T, U> &pair) {
  return os << '(' << pair.first << ", " << pair.second << ')';
}

template<class T>
ostream& operator<< (ostream& os, const vector<T> &arr) {
  os << '[';
  for (const T &v : arr) os << v << ", ";
  return os << ']';
}

// TODO replace this with my own linked list because I can't guarantee that
// the inserts will be constant time rahhh
struct UnionFind {
  vector<int> par;
  vector<int> id;
  vector<list<int>> next_queue;
  UnionFind(int n) : par(n, -1), id(n), next_queue(n) {
    iota(begin(id), end(id), 0);
  }
  int find(int i) {
    return (par[i] < 0) ? i : (par[i] = find(par[i]));
  }
  void join(int i, int j, int new_id) {
    i = find(i); j = find(j);
    id[i] = id[j] = new_id;
    if (i == j) return;
    if (par[i] > par[j]) swap(i, j);
    par[i] += par[j]; par[j] = i;
    next_queue[i].insert(next_queue[i].end(), next_queue[j].begin(), next_queue[j].end());
  }
  int rep(int i) {return id[find(i)];}
  list<int>& get_next_queue(int i) {
    return next_queue[find(i)];
  }
};

// return a partial partition of {0, ..., n-1} corresponding to sccs that cannot reach any other nodes
vector<vector<int>> partial_sccs(const vector<int> &keys, const vector<vector<pair<int, int>>> &graph) {
  int n = keys.size();

  vector<vector<int>> ans;
  // finalized[i] = can this vtx currently reach a vtx in ans?
  vector<bool> finalized(n);
  
  vector<int> cstack;
  vector<int> par(n, -1);
  vector<int> stack_pos(n, -1);
  vector<int> highest_reachable(n, n);
  vector<int> lowest_on_stack(n, -1);

  UnionFind current_sccs(n);

  vector<vector<pair<int, int>>> pending_color_queue(n);

  auto push_to_queue = [&] (int from, int to) {
      current_sccs.get_next_queue(from).push_back(to);
  };

  auto mark_vertex = [&] (int c) {
    for (auto &[from, to] : pending_color_queue[c]) {
      int source = current_sccs.rep(from);
      push_to_queue(source, to);
    }
    pending_color_queue[c].clear();
  };

  auto register_edge = [&] (int from, int to, int c) {
    pending_color_queue[c].emplace_back(from, to); 
  };
  
  function<bool(int, int)> visit; visit = [&] (int v, int p) {
    // cerr << "visit " << v << " from " << p << endl;
    if (stack_pos[v] != -1) return true;
    if (finalized[v]) return false;
    par[v] = p;
    stack_pos[v] = cstack.size();
    cstack.push_back(v);
    lowest_on_stack[keys[v]] = stack_pos[v]; 
    mark_vertex(keys[v]);
    // cerr << cstack << endl;

    vector<pair<int, int>> neighbors = graph[v];
    for (auto [to, c] : neighbors) register_edge(v, to, c);
    sort(begin(neighbors), end(neighbors), [&] (pair<int, int> n1, pair<int, int> n2) {
      return lowest_on_stack[n1.second] < lowest_on_stack[n2.second];
    });

    int highest_stack_pos = stack_pos[v];
    /*
    cerr << v << " neighbors:";
    for (auto [to, c] : neighbors) cerr << "(" << to << ", " << c << ", " << lowest_on_stack[c] << ") ";
    cerr << endl;
    */
    for (auto [to, c] : neighbors) if (c == keys[v]) current_sccs.get_next_queue(v).push_back(to);

    auto update_hsp = [&] (int hsp) {
      // cerr << v << "update hsp with " << hsp << endl;
      highest_stack_pos = min(highest_stack_pos, hsp);
      while (!neighbors.empty() && lowest_on_stack[neighbors.back().second] >= highest_stack_pos) {
        int to = neighbors.back().first;
        neighbors.pop_back();
        // cerr << "(hsp update) add " << v << "to queue" << endl;
        current_sccs.get_next_queue(v).push_back(to);
      }
    };

    while (!current_sccs.get_next_queue(v).empty()) {
      int cur = current_sccs.get_next_queue(v).back(); current_sccs.get_next_queue(v).pop_back(); 
      if (stack_pos[cur] != -1) {
        // cerr << "! back edge " << v << " " << cur << endl;
        if (finalized[cur]) return false;
        // this back edge gives us access to some colors. push them to our queue
        update_hsp(stack_pos[current_sccs.rep(cur)]);
        // merge sccs as needed to account for this back edge
        while (current_sccs.rep(v) != current_sccs.rep(cur)) {
          int parent = par[current_sccs.rep(v)];
          // cerr << "joining " << v << " with " << parent << endl;
          current_sccs.join(v, parent, current_sccs.rep(parent));
        }
      } else {
        if (!visit(cur, v)) return false;
        update_hsp(highest_reachable[v]);
        update_hsp(stack_pos[current_sccs.rep(v)]);
      }
    }
    // cerr << "visit " << v << " from " << p << " ending " <<  endl;
    if (v == current_sccs.rep(v) && highest_stack_pos <= stack_pos[v]) {
      // cerr << "! make new SCC" << endl;
      // new scc here
      ans.emplace_back(begin(cstack) + stack_pos[v], end(cstack));
      return false;
    }
    highest_reachable[v] = highest_stack_pos;
    return true; 
  };

  for (int i = 0; i < n; ++i) {
    if (finalized[i]) continue;
    visit(i, -1);
    for (int v : cstack) {
      finalized[v] = true;
      lowest_on_stack[keys[v]] = -1;
    }
    cstack.clear();
  }
  return ans;
}

vector<int> find_reachable(vector<int> r, vector<int> u, vector<int> v, vector<int> c) {
  int n = r.size(), m = u.size();
  vector<vector<pair<int, int>>> graph(n);
  for (int i = 0; i < m; ++i) {
    graph[u[i]].emplace_back(v[i], c[i]);
    graph[v[i]].emplace_back(u[i], c[i]);
  }
  vector<vector<int>> sccs = partial_sccs(r, graph);
  int min_scc_size = n;
  for (auto &scc : sccs) min_scc_size = min((int)scc.size(), min_scc_size);
  vector<int> ans(n);
  for (auto &scc : sccs) {
    if (scc.size() == min_scc_size) {
      for (int v : scc) ans[v] = 1;
    }
  }
  return ans;
}

Compilation message (stderr)

keys.cpp: In function 'std::vector<int> find_reachable(std::vector<int>, std::vector<int>, std::vector<int>, std::vector<int>)':
keys.cpp:166:20: warning: comparison of integer expressions of different signedness: 'std::vector<int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
  166 |     if (scc.size() == min_scc_size) {
      |         ~~~~~~~~~~~^~~~~~~~~~~~~~~

• Subtask 1: 9 / 9
• Subtask 2: 11 / 11
• Subtask 3: 17 / 17
• Subtask 4: 30 / 30
• Subtask 5: 33 / 33