IOI '18 - Doll (53pts)

#include "doll.h"
#include <set>
#include <tuple>
#include <cassert>
#include <iostream>

using namespace std;

const int INF = 1e9;

bool drain(int i, int root, vector<int> &X, vector<int> &Y) {
  bool isDrain = false;
  if (X[-i] < 0) {if (drain(X[-i], root, X, Y)) isDrain = true;}
  else if (X[-i] == 0 || X[-i] == INF) X[-i] = root, isDrain = true;
  if (Y[-i] < 0) {if (drain(Y[-i], root, X, Y)) isDrain = true;}
  else if (Y[-i] == 0 || Y[-i] == INF) Y[-i] = root, isDrain = true;
  return isDrain;
}

void create_circuit(int M, vector<int> A) {
  int N = A.size();
  vector<vector<int>> nxt(1+M);
  nxt[0].push_back(A[0]);
  for (int i = 0; i < N-1; ++i) {
    nxt[A[i]].push_back(A[i+1]);
  }
  nxt[A[N-1]].push_back(0);

  int S = 0;
  vector<int> X(1), Y(1);
  vector<int> C(1+M); // root switch for each trigger

  for (int i = 0; i <= M; ++i) {
    if (nxt[i].empty()) { // unused trigger
      C[i] = i;
      continue;
    }
    int p2 = 1, expo = 0;
    while (p2 < (int)nxt[i].size()) ++expo, p2 <<= 1;

    // Order targets (stored in nxt) for tree of switches
    vector<int> orderedNxt(1, 0);
    for (int j = 1, previousPow2 = 1; j <= expo; ++j, previousPow2 <<= 1) {
      vector<int> nextLevel(2*orderedNxt.size());
      for (int k = 0; k < (int)nextLevel.size(); ++k) {
        nextLevel[k] = orderedNxt[k>>1] + (k&1)*previousPow2;
      }
      orderedNxt = nextLevel;
    }
    for (int j = 0; j < p2; ++j) {
      if (orderedNxt[j] >= (int)nxt[i].size()) orderedNxt[j] = INF;
      else orderedNxt[j] = nxt[i][orderedNxt[j]];
    }

    // Now build switches for pairs of targets
    vector<int> curSwitches = orderedNxt;
    while (curSwitches.size() > 1) {
      vector<int> nextSwitches((int)curSwitches.size()>>1);
      for (int j = 0; j < (int)nextSwitches.size(); ++j) {
        int x = curSwitches[2*j], y = curSwitches[2*j+1];
        if (x==INF && y==INF) {// this should never happen
          nextSwitches[j] = INF;
          continue;
        }
        // create switch

        // Create new switch
        ++S;
        X.push_back(x);
        Y.push_back(y);
        nextSwitches[j] = -S;
      }
      curSwitches = nextSwitches;
    }

    // now curSwitches[0] is the root switch for this trigger
    C[i] = curSwitches[0];
  }

  // Now we need to "drain" the switches at the end so they all have state X at the end
  //  find the roots of the trees that need draining
  vector<int> needDrain;
  for (int i = 0; i <= M; ++i) {
    if (C[i] >= 0) {
      if (C[i] == 0 || C[i] == INF) C[i] = i;
    } else if (drain(C[i], C[i], X, Y) && i != A[N-1]) {
      needDrain.push_back(i);
    }
  }

  auto getBottom = [&C, &Y](int i) {
    int root = C[i]; i = C[i];
    while (true) {
      if (i > 0) {
        if (C[i] == root) return i;
        i = C[i];
      } else {
        if (Y[-i] == root) return i;
        i = Y[-i];
      }
    }
  };

  //  link up the drain chain drain chain
  int prevBottom = getBottom(A[N-1]);
  for (int i : needDrain) {
    // connect bottom (end) of previous to root of current
    if (prevBottom > 0) C[prevBottom] = C[i];
    else Y[-prevBottom] = C[i];
    // find bottom of current
    prevBottom = getBottom(i);
  }

  //   connect last bottom back to the origin
  if (prevBottom > 0) C[prevBottom] = 0;
  else Y[-prevBottom] = 0;

  X.erase(X.begin());
  Y.erase(Y.begin());
  answer(C, X, Y);
}