Codeforces Marathon Round 2 yarrr's final submit

cf_marathon_r2.cpp

#pragma GCC optimize("Ofast,no-stack-protector,tree-vectorize")
#pragma GCC target("avx,sse4,tune=native")

#include <iostream>
#include <cassert>
#include <string>
#include <vector>
#include <array>
#include <numeric>
#include <algorithm>
#include <unordered_map>
#include <unordered_set>
#include <tuple>

static constexpr int kDifferentColors = 10;
static constexpr int kBottlesInHands = 10;
static constexpr int kAnimalsCount = 10;

static constexpr int kBestStatesToHold = 650;
static constexpr double kScoreDecreaseCoefficient = 1.07;
static constexpr bool kDebug = false;
static constexpr double kEps = 1e-9;

using Positions = std::array<int, kAnimalsCount>;

class Input {
public:
  void read(std::istream& in) {
    int differentColors, bottlesInHands, animalsCount;
    in >> tapeLength_ >> bottlesCount_ >> differentColors >> bottlesInHands >> animalsCount;
    assert(differentColors == kDifferentColors);
    assert(bottlesInHands == kBottlesInHands);
    assert(animalsCount == kAnimalsCount);
    bottlesCount_ += bottlesInHands;
    in >> tape_;
    assert(tape_.size() == tapeLength_);
    in >> bottles_;
    assert(bottles_.size() == bottlesCount_);
  }

  int tapeLength() const {
    return tapeLength_;
  }

  int bottlesCount() const {
    return bottlesCount_;
  }

  const std::string& tape() const {
    return tape_;
  }

  const std::string& bottles() const {
    return bottles_;
  }

private:
  int tapeLength_;
  int bottlesCount_;
  std::string tape_;
  std::string bottles_;
};

struct CertificateAction {
  int animalId;
  char bottleColor;
};

class FastHashSet {
public:
  static constexpr size_t kSize = 32768;

  FastHashSet() : keys_(kSize, 0) {}

  void clear() {
    std::fill(keys_.begin(), keys_.end(), 0);
  }

  bool insert(uint64_t key) {
    assert(key != 0);
    uint32_t hash = key & (kSize - 1);
    while (keys_[hash] != 0 && keys_[hash] != key) {
      ++hash;
      if (hash == kSize) {
        hash = 0;
      }
    }
    if (keys_[hash] == 0) {
      keys_[hash] = key;
      return true;
    } else {
      return false;
    }
  }

private:
  std::vector<uint64_t> keys_;
};

class Answer {
public:
  Answer(std::vector<CertificateAction>&& certificate) : certificate_(std::move(certificate)) {}

  void print(std::ostream& out) {
    for (const auto& action : certificate_) {
      out << action.animalId << " " << action.bottleColor << "\n";
    }
    out.flush();
  }

private:
  std::vector<CertificateAction> certificate_;
};

struct GameState {
  Positions animalPositions;
  Positions bottlesInHands;
  int nextBottleId;

  GameState() {
    std::iota(animalPositions.begin(), animalPositions.end(), 0);
    std::iota(bottlesInHands.begin(), bottlesInHands.end(), 0);
    nextBottleId = bottlesInHands.size();
  }
};

template <typename T>
struct PairWithStateId {
  int stateId;
  T data;
};

using CertificateActionWithStateId = PairWithStateId<CertificateAction>;
using GameStateWithStateId = PairWithStateId<GameState>;

struct Candidate {
  double score;
  int previousStateId;
  const GameState* previousState;
  int animalId;
  int bottleId;
};

class Solution {
public:
  Solution(const Input& input) :
      input_(input),
      totalIterations_(input_.bottlesCount() - kBottlesInHands),
      lastStateId_(0)
  {
    precalcScoreCoefficients();
    precalcInstantColorJumps();
    certificateHistory_.resize(totalIterations_ * kBestStatesToHold);
    bestNextCandidates_.reserve(2 * kBestStatesToHold);
    currentStates_.reserve(kBestStatesToHold);
    nextStates_.reserve(kBestStatesToHold);
    if (kDebug) {
      allPositions_.resize(totalIterations_ * kBestStatesToHold);
      bestStates_.resize(totalIterations_);
    }
  }

  Answer solve() {
    currentStates_.emplace_back(GameStateWithStateId{nextStateId(), GameState()});

    if (kDebug) {
      allPositions_[currentStates_.front().stateId] = currentStates_.front().data.animalPositions;
    }

    for (int iteration = 0; iteration < totalIterations_; iteration++) {
      double minScoreToAdd = 0.0;

      usedScores_.clear();
      bestNextCandidates_.clear();

      if (kDebug) {
        bestStates_[iteration] = currentStates_.front().stateId;
      }

      auto removeWorstNextStates = [&]() {
        if (bestNextCandidates_.size() >= kBestStatesToHold) {
          std::nth_element(
              bestNextCandidates_.begin(),
              bestNextCandidates_.begin() + kBestStatesToHold,
              bestNextCandidates_.end(),
              [&](const Candidate& lhs, const Candidate& rhs) {
                return lhs.score > rhs.score;
              });
          bestNextCandidates_.resize(kBestStatesToHold);
          minScoreToAdd = bestNextCandidates_.front().score;
          for (const auto& state : bestNextCandidates_) {
            minScoreToAdd = std::min(minScoreToAdd, state.score);
          }
        }
      };

      for (const auto& currentState : currentStates_) {
        Positions sortedAnimalPositions = currentState.data.animalPositions;
        std::sort(sortedAnimalPositions.begin(), sortedAnimalPositions.end());
        for (int animalId = 0; animalId < kAnimalsCount; animalId++) {
          std::array<bool, kDifferentColors> colorUsed;
          colorUsed.fill(false);
          int animalPosition = currentState.data.animalPositions[animalId];
          int animalCursor = std::distance(
              sortedAnimalPositions.begin(),
              std::lower_bound(sortedAnimalPositions.begin(), sortedAnimalPositions.end(), animalPosition)
          );
          int minFinalAnimalPosition = 0;
          for (int bottleId = 0; bottleId < kBottlesInHands; bottleId++) {
            int bottlePosition = currentState.data.bottlesInHands[bottleId];
            assert(bottlePosition < input_.bottlesCount());
            char bottleColor = input_.bottles()[bottlePosition];
            int bottleColorId = bottleColor - 'A';
            if (colorUsed[bottleColorId]) {
              continue;
            }
            colorUsed[bottleColorId] = true;
            int finalAnimalPosition = 0;
            auto newScore = movePositionsAndCalcScore(
                sortedAnimalPositions,
                animalPosition,
                animalCursor,
                bottleColorId,
                finalAnimalPosition,
                minFinalAnimalPosition
            );
            if (newScore <= minScoreToAdd + kEps) {
              minFinalAnimalPosition = std::max(minFinalAnimalPosition, finalAnimalPosition);
              continue;
            }
            uint64_t newScoreHash = *reinterpret_cast<const uint64_t*>(&newScore);
            bool inserted = usedScores_.insert(newScoreHash);
            if (!inserted) {
              continue;
            }
            bestNextCandidates_.push_back(
                Candidate{newScore, currentState.stateId, &currentState.data, animalId, bottleId}
            );
            if (bestNextCandidates_.size() >= 2 * kBestStatesToHold) {
              removeWorstNextStates();
            }
          }
        }
      }

      removeWorstNextStates();
      std::sort(bestNextCandidates_.begin(), bestNextCandidates_.end(), [&](const Candidate& lhs, const Candidate& rhs) {
        return lhs.score > rhs.score;
      });
      nextStates_.clear();

      for (const auto& candidate : bestNextCandidates_) {
        GameState newState = moveState(
            *candidate.previousState,
            candidate.animalId,
            candidate.bottleId
        );
        int newStateId = nextStateId();
        nextStates_.emplace_back(GameStateWithStateId{newStateId, newState});
        int bottlePosition = (*candidate.previousState).bottlesInHands[candidate.bottleId];
        assert(bottlePosition < input_.bottlesCount());
        char bottleColor = input_.bottles()[bottlePosition];
        CertificateAction action{candidate.animalId, bottleColor};
        CertificateActionWithStateId actionWithStateId{candidate.previousStateId, action};
        certificateHistory_[newStateId] = actionWithStateId;
        if (kDebug) {
          allPositions_[newStateId] = newState.animalPositions;
          auto checkPositions = newState.animalPositions;
          std::sort(checkPositions.begin(), checkPositions.end());
          assert(std::abs(evaluateAnimalPositionsScore(checkPositions) - candidate.score) < kEps);
        }
      }

      currentStates_.swap(nextStates_);
    }


    auto finalState = getFinalState();

    std::cerr << "Min animal position at finish: " << finalScore(finalState.data) << std::endl;

    auto certificate = buildCertificate(finalState.stateId);

    if (kDebug) {
      auto checkGameState = simulate(certificate);
      std::cerr << "Min animal position at finish for simulation: " << finalScore(checkGameState) << std::endl;
      assert(checkGameState.animalPositions == finalState.data.animalPositions);
    }

    return Answer(std::move(certificate));
  }

private:
  inline int nextStateId() {
    return lastStateId_++;
  }

  void precalcScoreCoefficients() {
    double k = 1.0;
    for (int i = 0; i < kAnimalsCount + kAnimalsCount; i++) {
      kScoreCoefficients_[i] = k;
      k /= kScoreDecreaseCoefficient;
    }
  }

  void precalcInstantColorJumps() {
    for (int colorId = 0; colorId < kDifferentColors; colorId++) {
      int tapeLength = input_.tapeLength();
      auto& jumps = instantColorJumps_[colorId];
      jumps.resize(tapeLength);
      int lastColorPosition = tapeLength;
      for (int position = tapeLength - 1; position >= 0; position--) {
        jumps[position] = lastColorPosition;

        int colorIdOnPosition = input_.tape()[position] - 'A';
        if (colorIdOnPosition == colorId) {
          lastColorPosition = position;
        }
      }
    }
  }

  inline double movePositionsAndCalcScore(
      const Positions& animalPositions,
      int animalPosition,
      int startAnimalCursor,
      int colorId,
      int& finalAnimalPosition,
      int minFinalAnimalPosition
  ) const {
    int finishAnimalCursor = startAnimalCursor;

    while (animalPosition < input_.tapeLength()) {
      animalPosition = instantColorJumps_[colorId][animalPosition];

      while (finishAnimalCursor < animalPositions.size() &&
          animalPositions[finishAnimalCursor] < animalPosition) {
        ++finishAnimalCursor;
      }

      if (finishAnimalCursor == animalPositions.size() ||
          animalPositions[finishAnimalCursor] != animalPosition) {
        break;
      }
    }

    if (animalPosition <= minFinalAnimalPosition) {
      return 0.0;
    }

    finalAnimalPosition = animalPosition;

    if (finishAnimalCursor > startAnimalCursor) {
      return trickyEvaluateAnimalPositionsScore(
          animalPositions,
          startAnimalCursor,
          finishAnimalCursor,
          animalPosition
      );
    } else {
      assert(animalPosition == input_.tapeLength());
      assert(animalPositions[startAnimalCursor] == input_.tapeLength());
      return evaluateAnimalPositionsScore(animalPositions);
    }
  }

  GameState moveState(GameState currentState, int animalId, int bottleId) const {
    auto& animalPosition = currentState.animalPositions[animalId];
    auto& bottlePosition = currentState.bottlesInHands[bottleId];
    assert(bottlePosition < input_.bottlesCount());
    char bottleColorId = input_.bottles()[bottlePosition] - 'A';
    bool over = false;

    while (animalPosition < input_.tapeLength() && !over) {
      animalPosition = instantColorJumps_[bottleColorId][animalPosition];
      if (animalPosition >= input_.tapeLength()) {
        animalPosition = input_.tapeLength();
        over = true;
      } else {
        bool canSkip = false;
        for (int otherAnimalId = 0; otherAnimalId < kAnimalsCount; otherAnimalId++) {
          if (otherAnimalId != animalId &&
              currentState.animalPositions[otherAnimalId] == animalPosition) {
            canSkip = true;
          }
        }
        if (!canSkip) {
          over = true;
        }
      }
    }

    assert(currentState.nextBottleId < input_.bottlesCount());
    bottlePosition = currentState.nextBottleId++;

    return currentState;
  }

  inline double evaluateAnimalPositionsScore(
      const Positions& animalPositions
  ) const {
    double sum = 0;
    for (int i = 0; i < kAnimalsCount; i++) {
      sum += animalPositions[i] * kScoreCoefficients_[i];
    }
    return sum;
  }

  inline double trickyEvaluateAnimalPositionsScore(
      const Positions& animalPositions,
      int startAnimalCursor,
      int finishAnimalCursor,
      int animalPosition
  ) const {
    assert(startAnimalCursor < finishAnimalCursor);
    --finishAnimalCursor;
    double sum = 0.0;
    for (int i = 0; i < kAnimalsCount; i++) {
      if (i >= startAnimalCursor && i < finishAnimalCursor) {
        sum += animalPositions[i + 1] * kScoreCoefficients_[i];
      } else if (i == finishAnimalCursor) {
        sum += animalPosition * kScoreCoefficients_[i];
      } else {
        sum += animalPositions[i] * kScoreCoefficients_[i];
      }
    }
    return sum;
  }

  inline int finalScore(const GameState& state) const {
    return *std::min_element(
        state.animalPositions.begin(),
        state.animalPositions.end()
    );
  }

  GameStateWithStateId getFinalState() const {
    GameStateWithStateId bestFinalState;
    for (const auto& currentState : currentStates_) {
      if (finalScore(currentState.data) > finalScore(bestFinalState.data)) {
        bestFinalState = currentState;
      }
    }
    return bestFinalState;
  }

  std::vector<CertificateAction> buildCertificate(int finalStateId) const {
    std::vector<CertificateAction> certificate;
    std::vector<int> states;
    certificate.reserve(totalIterations_);
    if (kDebug) {
      states.reserve(totalIterations_);
    }
    int stateCursor = finalStateId;
    while (stateCursor != 0) {
      const auto& action = certificateHistory_[stateCursor];
      certificate.push_back(action.data);
      stateCursor = action.stateId;
      if (kDebug) {
        states.push_back(stateCursor);
      }
    }
    std::reverse(certificate.begin(), certificate.end());
    if (kDebug) {
      auto printArray = [&](Positions positions) {
        std::cerr << "[";
        for (int i = 0; i < positions.size(); i++) {
          std::cerr << positions[i];
          std::cerr << ":" << input_.tape()[positions[i]];
          if (i + 1 != positions.size()) {
            std::cerr << ", ";
          }
        }
        std::cerr << "]";
      };
      std::reverse(states.begin(), states.end());
      assert(states.size() == totalIterations_);
      for (int iteration = 0; iteration < totalIterations_; iteration++) {
        auto bestPositions = allPositions_[bestStates_[iteration]];
        auto currentPositions = allPositions_[states[iteration]];
        std::sort(bestPositions.begin(), bestPositions.end());
        std::sort(currentPositions.begin(), currentPositions.end());
        double bestScore = evaluateAnimalPositionsScore(bestPositions);
        double currentScore = evaluateAnimalPositionsScore(currentPositions);
        std::cerr << "iteration: " << (iteration + 1);
        std::cerr << "\t best local (" << bestScore << "): ";
        printArray(bestPositions);
        std::cerr << "\tbest optimal (" << currentScore << "): ";
        printArray(currentPositions);
        std::cerr << "\tdistance: " << (states[iteration] - bestStates_[iteration]);
        std::cerr << std::endl;
        assert(bestScore >= currentScore);
      }
    }
    return certificate;
  }

  GameState simulate(const std::vector<CertificateAction>& certificate) {
    GameState gameState;
    for (auto action : certificate) {
      auto& animalPosition = gameState.animalPositions[action.animalId];
      bool over = false;
      while (!over) {
        ++animalPosition;
        if (animalPosition >= input_.tapeLength()) {
          animalPosition = input_.tapeLength();
          over = true;
        } else {
          char bottleColor = input_.tape()[animalPosition];
          if (bottleColor == action.bottleColor) {
            bool canSkip = false;
            for (int otherAnimalId = 0; otherAnimalId < kAnimalsCount; otherAnimalId++) {
              if (gameState.animalPositions[otherAnimalId] == animalPosition &&
                  otherAnimalId != action.animalId) {
                canSkip = true;
              }
            }
            if (!canSkip) {
              over = true;
            }
          }
        }
      }

      bool bottleFound = false;
      for (auto& bottlePosition : gameState.bottlesInHands) {
        if (input_.bottles()[bottlePosition] == action.bottleColor) {
          bottlePosition = gameState.nextBottleId++;
          bottleFound = true;
          break;
        }
      }
      assert(bottleFound);
    }
    return gameState;
  }

  const Input& input_;
  const int totalIterations_;
  int lastStateId_;
  std::vector<CertificateActionWithStateId> certificateHistory_;
  std::array<std::vector<int>, kDifferentColors> instantColorJumps_;
  std::array<double, 2 * kAnimalsCount> kScoreCoefficients_;
  FastHashSet usedScores_;
  std::vector<Candidate> bestNextCandidates_;
  std::vector<GameStateWithStateId> currentStates_;
  std::vector<GameStateWithStateId> nextStates_;
  std::vector<Positions> allPositions_;
  std::vector<int> bestStates_;
};

int main() {
  Input input;
  input.read(std::cin);
  Solution solution(input);
  Answer answer = solution.solve();
  answer.print(std::cout);
  return 0;
}