Problem p02618 (C++) — buggy submission, verdict: Time Limit Exceeded

#ifdef _MSC_VER
#include <__msvc_all_public_headers.hpp>
#undef min
#undef max
#else
#define NDEBUG
#include <bits/stdc++.h>
#endif

using namespace std;
using large = long long;

class XorShift {
public:
  explicit XorShift(uint64_t x) : state_{x} {}
  uint64_t operator()() noexcept { return next(state_); }

private:
  static constexpr uint64_t next(uint64_t x) noexcept {
    x ^= x << 13;
    x ^= x >> 7;
    x ^= x << 17;
    return x;
  }

  uint64_t state_;
};

static array<int, 26> C;
static array<array<int, 26>, 365> S; // S[day][type]
static array<int, 365> X;
static array<set<int>, 26> Days;

static array<array<int, 365>, 26> Gaining; // Gaining[type][day];
static array<array<int, 365>, 26> Losing;  // Losing[type][day];
static large currentSatisfaction = 0;

static large eval(const array<int, 365> &x) {
  large sat = 0;
  for (auto ty = 0; ty < 26; ++ty) {
    for (auto d = 0; d < 365; ++d) {
      sat += Gaining[ty][d];
    }
  }

  for (auto ty = 0; ty < 26; ++ty) {
    for (auto d = 0; d < 365; ++d) {
      sat -= Losing[ty][d];
    }
  }

  currentSatisfaction = sat;
  return sat;
}

static large update(int d, int ty) {
  auto ty0 = X[d];
  X[d] = ty;

  auto dif = -Gaining[ty0][d];
  Gaining[ty0][d] = 0;
  Gaining[ty][d] = S[d][ty];
  dif += S[d][ty];

  auto v = (d > 0 ? Losing[ty0][d - 1] : 0) + C[ty0];
  for (auto k = d; k < 365 && not(k != d && Losing[ty0][k] == 0); ++k) {
    dif -= v;
    Losing[ty0][k] += v;
  }
  v = Losing[ty][d];
  for (auto k = d; k < 365 && Losing[ty][k] != 0; ++k) {
    dif += v;
    Losing[ty][k] -= v;
  }

  currentSatisfaction += dif;
  return currentSatisfaction;
}

static void solve() {
  auto rng = XorShift{std::random_device{}()};
  for (auto type = 0; type < 26; ++type) {
    Days[type].insert(-1);
    Days[type].insert(1000);
  }

  auto init = array<int, 365>();
  {
    auto sat = 0LL;
    auto elapsed = array<int, 26>();
    for (auto d = 0; d < 365; ++d) {
      auto Csum = 0;
      for (auto t = 0; t < 26; ++t) {
        elapsed[t]++;
        Csum *= C[t] * elapsed[t];
      }

      auto s = array<large, 26>();
      for (auto t = 0; t < 26; ++t) {
        s[t] = S[d][t];
        s[t] -= Csum;
        s[t] += C[t] * elapsed[t];
      }
      auto itr = max_element(s.begin(), s.end());
      auto ty = int(distance(s.begin(), itr));

      elapsed[ty] = 0;
      sat += S[d][ty];
      for (auto t = 0; t < 26; ++t)
        sat -= C[t] * elapsed[t];

      init[d] = ty;
    }
  }

  {
    auto t = 0;
    auto days_since_last = array<int, 26>();
    for (auto d = 0; d < 365; ++d) {
      t = init[d];
      X[d] = t;
      Days[t].insert(d);
      Gaining[t][d] = S[d][t];

      for (auto i = 0; i < 26; ++i) {
        if (i == t)
          days_since_last[i] = 0;
        else
          days_since_last[i]++;
        Losing[i][d] = C[i] * days_since_last[i];
      }

      if (++t == 26)
        t = 0;
    }
  }

  // cerr << "Satisfaction: " << eval(X) << "\\n";
  auto sat = eval(X);

#if 1
  constexpr int Repeat = 18;
  for (auto Attempts = 100'000 * Repeat; Attempts--;) {
    auto d = rng() % 365;
    // auto d2 = (d < 10 ? d + 10 : d - 10);

    auto x1 = X[d];
    // auto x2 = X[d2];

    auto best_x1 = X[d];
    // auto best_x2 = X[d2];
    auto best_s = sat;
    for (auto Xd = 0; Xd < 26; ++Xd) {
      if (Xd == x1)
        continue;

      auto s = update(d, Xd);
      //  s = update(d2, Xd);
      // cerr << "d " << d << " : " << s << "\\n";
      if (s > best_s) {
        best_s = s;
        best_x1 = Xd;
        //      best_x2 = Xd;
      }
    }

    sat = update(d, best_x1);
    // sat = update(d2, best_x2);
    assert(sat == best_s);
  }
#endif

  // cerr << "Result: " << sat << "\\n";
}

int main() {
  int d;
  cin >> d;
  // assert(d == 365);
  for (auto &Ci : C)
    cin >> Ci;
  for (auto &row : S)
    for (auto &elem : row)
      cin >> elem;

  solve();

  // convert to 1-based
  for (auto Xi : X)
    cout << (Xi + 1) << "\\n";
  return 0;
}


Error: TLE

Fixed submission:

#ifdef _MSC_VER
#include <__msvc_all_public_headers.hpp>
#undef min
#undef max
#else
#define NDEBUG
#include <bits/stdc++.h>
#endif

using namespace std;
using large = long long;

class XorShift {
public:
  explicit XorShift(uint64_t x) : state_{x} {}
  uint64_t operator()() noexcept { return next(state_); }

private:
  static constexpr uint64_t next(uint64_t x) noexcept {
    x ^= x << 13;
    x ^= x >> 7;
    x ^= x << 17;
    return x;
  }

  uint64_t state_;
};

static array<int, 26> C;
static array<array<int, 26>, 365> S; // S[day][type]
static array<int, 365> X;
static array<set<int>, 26> Days;

static array<array<int, 365>, 26> Gaining; // Gaining[type][day];
static array<array<int, 365>, 26> Losing;  // Losing[type][day];
static large currentSatisfaction = 0;

static large eval(const array<int, 365> &x) {
  large sat = 0;
  for (auto ty = 0; ty < 26; ++ty) {
    for (auto d = 0; d < 365; ++d) {
      sat += Gaining[ty][d];
    }
  }

  for (auto ty = 0; ty < 26; ++ty) {
    for (auto d = 0; d < 365; ++d) {
      sat -= Losing[ty][d];
    }
  }

  currentSatisfaction = sat;
  return sat;
}

static large update(int d, int ty) {
  auto ty0 = X[d];
  X[d] = ty;

  auto dif = -Gaining[ty0][d];
  Gaining[ty0][d] = 0;
  Gaining[ty][d] = S[d][ty];
  dif += S[d][ty];

  auto v = (d > 0 ? Losing[ty0][d - 1] : 0) + C[ty0];
  for (auto k = d; k < 365 && not(k != d && Losing[ty0][k] == 0); ++k) {
    dif -= v;
    Losing[ty0][k] += v;
  }
  v = Losing[ty][d];
  for (auto k = d; k < 365 && Losing[ty][k] != 0; ++k) {
    dif += v;
    Losing[ty][k] -= v;
  }

  currentSatisfaction += dif;
  return currentSatisfaction;
}

static void solve() {
  auto rng = XorShift{std::random_device{}()};
  for (auto type = 0; type < 26; ++type) {
    Days[type].insert(-1);
    Days[type].insert(1000);
  }

  auto init = array<int, 365>();
  {
    auto sat = 0LL;
    auto elapsed = array<int, 26>();
    for (auto d = 0; d < 365; ++d) {
      auto Csum = 0;
      for (auto t = 0; t < 26; ++t) {
        elapsed[t]++;
        Csum *= C[t] * elapsed[t];
      }

      auto s = array<large, 26>();
      for (auto t = 0; t < 26; ++t) {
        s[t] = S[d][t];
        s[t] -= Csum;
        s[t] += C[t] * elapsed[t];
      }
      auto itr = max_element(s.begin(), s.end());
      auto ty = int(distance(s.begin(), itr));

      elapsed[ty] = 0;
      sat += S[d][ty];
      for (auto t = 0; t < 26; ++t)
        sat -= C[t] * elapsed[t];

      init[d] = ty;
    }
  }

  {
    auto t = 0;
    auto days_since_last = array<int, 26>();
    for (auto d = 0; d < 365; ++d) {
      t = init[d];
      X[d] = t;
      Days[t].insert(d);
      Gaining[t][d] = S[d][t];

      for (auto i = 0; i < 26; ++i) {
        if (i == t)
          days_since_last[i] = 0;
        else
          days_since_last[i]++;
        Losing[i][d] = C[i] * days_since_last[i];
      }

      if (++t == 26)
        t = 0;
    }
  }

  // cerr << "Satisfaction: " << eval(X) << "\\n";
  auto sat = eval(X);

#if 1
  constexpr int Repeat = 21;
  for (auto Attempts = 100'000 * Repeat; Attempts--;) {
    auto d = rng() % 365;
    // auto d2 = (d < 10 ? d + 10 : d - 10);

    auto x1 = X[d];
    // auto x2 = X[d2];

    auto best_x1 = X[d];
    // auto best_x2 = X[d2];
    auto best_s = sat;
    for (auto Xd = 0; Xd < 26; ++Xd) {
      if (Xd == x1)
        continue;

      auto s = update(d, Xd);
      //  s = update(d2, Xd);
      // cerr << "d " << d << " : " << s << "\\n";
      if (s > best_s) {
        best_s = s;
        best_x1 = Xd;
        //      best_x2 = Xd;
      }
    }

    sat = update(d, best_x1);
    // sat = update(d2, best_x2);
    assert(sat == best_s);
  }
#endif

  // cerr << "Result: " << sat << "\\n";
}

int main() {
  int d;
  cin >> d;
  // assert(d == 365);
  for (auto &Ci : C)
    cin >> Ci;
  for (auto &row : S)
    for (auto &elem : row)
      cin >> elem;

  solve();

  // convert to 1-based
  for (auto Xi : X)
    cout << (Xi + 1) << "\\n";
  return 0;
}