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

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using ull = unsigned long long;
using ld = long double;
using pii = pair<int, int>;
using pll = pair<ll, ll>;
using pli = pair<ll, int>;
using pil = pair<int, ll>;
using uint = unsigned int;
template <typename T> using Graph = vector<vector<T>>;
const int MOD = 1e9 + 7;
const ld PI = acosl(-1);

struct timer {
  std::chrono::system_clock::time_point start, end;

  timer() { start = std::chrono::system_clock::now(); }

  double get() {
    end = std::chrono::system_clock::now();
    return std::chrono::duration_cast<std::chrono::milliseconds>(end - start)
        .count();
  }
};

int D;
int c[26], s[365][26], t[365], last[26], ans[366];

void calc() {
  for (int i = 0; i < D; ++i) {
    ans[i + 1] = ans[i] + s[i][t[i]];
    last[t[i]] = i + 1;
    for (int j = 0; j < 26; ++j)
      ans[i + 1] -= c[j] * (i + 1 < last[j] ? i + 1 : i + 1 - last[j]);
  }
}

void change(int d, int q) {
  int old = t[d];
  if (old == q)
    return;

  t[d] = q, last[old] = 0, last[q] = 0;
  int dif = 0, pos1 = 0, pos2 = 0;
  for (int i = 0; i < D; ++i) {
    last[t[i]] = i + 1;
    if (i < d)
      pos1 = last[q], pos2 = last[old];
    else {
      ans[i + 1] += s[d][q] - s[d][old];
      if (last[q] == d + 1)
        dif += c[q] * (d + 1 - pos1);
      if (last[old] <= d + 1)
        dif -= c[old] * (d + 1 - pos2);
    }
    ans[i + 1] += dif;
  }
}

void greedy() {
  for (int i = 0; i < D - 2; ++i) {
    ans[i + 3] = INT_MIN;
    int maidx1 = 0, maidx2 = 0, maidx3 = 0;
    for (int k1 = 0; k1 < 26; ++k1) {
      int old1 = last[k1];
      for (int k2 = 0; k2 < 26; ++k2) {
        int old2 = last[k2];
        for (int k3 = 0; k3 < 26; ++k3) {
          int old3 = last[k3];
          int res = ans[i] + s[i][k1] + s[i + 1][k2] + s[i + 2][k3], res1, res2;
          last[k1] = i + 1;
          for (int j = 0; j < 26; ++j)
            res -= c[j] * (i + 1 - last[j]);
          res1 = res;
          last[k2] = i + 2;
          for (int j = 0; j < 26; ++j)
            res -= c[j] * (i + 2 - last[j]);
          res2 = res;
          last[k3] = i + 3;
          for (int j = 0; j < 26; ++j)
            res -= c[j] * (i + 3 - last[j]);

          if (ans[i + 3] < res) {
            ans[i + 3] = res;
            ans[i + 2] = res2 - s[i + 2][k3];
            ans[i + 1] = res1 - s[i + 2][k3] - s[i + 1][k2];
            maidx1 = k1;
            maidx2 = k2;
            maidx3 = k3;
          }
          last[k3] = old3;
          last[k2] = old2;
          last[k1] = old1;
        }
      }
    }
    t[i] = maidx1;
    last[t[i]] = i + 1;
    if (i == D - 3) {
      t[i + 1] = maidx2;
      t[i + 2] = maidx3;
      last[t[i + 1]] = i + 2;
      last[t[i + 2]] = i + 3;
    }
  }
}

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);

  timer tm;

  cin >> D;
  for (int i = 0; i < 26; ++i)
    cin >> c[i];
  for (int i = 0; i < D; ++i)
    for (int j = 0; j < 26; ++j)
      cin >> s[i][j];

  greedy();
  std::mt19937_64 gen(1);
  std::uniform_int_distribution<> dist1(1, D);
  std::uniform_int_distribution<> dist2(1, 26);

  while (tm.get() < 1900) {
    stack<pii> stk;
    int pres = ans[D];
    for (int j = 0; j < 40; ++j) {
      int d = dist1(gen), q = dist2(gen);
      d--, q--;
      stk.emplace(pii{d, t[d]});
      change(d, q);
      if (ans[D] >= pres)
        break;
    }
    if (ans[D] < pres) {
      while (!stk.empty()) {
        auto [d, q] = stk.top();
        change(d, q);
        stk.pop();
      }
    }
  }
  for (int i = 0; i < D; ++i)
    cout << t[i] + 1 << '\\n';

  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using ull = unsigned long long;
using ld = long double;
using pii = pair<int, int>;
using pll = pair<ll, ll>;
using pli = pair<ll, int>;
using pil = pair<int, ll>;
using uint = unsigned int;
template <typename T> using Graph = vector<vector<T>>;
const int MOD = 1e9 + 7;
const ld PI = acosl(-1);

struct timer {
  std::chrono::system_clock::time_point start, end;

  timer() { start = std::chrono::system_clock::now(); }

  double get() {
    end = std::chrono::system_clock::now();
    return std::chrono::duration_cast<std::chrono::milliseconds>(end - start)
        .count();
  }
};

int D;
int c[26], s[365][26], t[365], ans[366];
vector<set<int>> last(26);

void calc() {
  for (int i = 0; i < D; ++i) {
    ans[i + 1] = ans[i] + s[i][t[i]];
    last[t[i]].emplace(i + 1);
    for (int j = 0; j < 26; ++j) {
      int tmp = *last[j].rbegin();
      if (i + 1 < tmp)
        ans[i + 1] -= c[j] * (i + 1);
      else
        ans[i + 1] -= c[j] * (i + 1 - tmp);
    }
  }
}

void change(int d, int q) {
  int old = t[d];
  if (old == q)
    return;

  t[d] = q;
  last[old].erase(d + 1);
  auto itr1 = last[q].lower_bound(d + 1);
  auto itr2 = last[old].lower_bound(d + 1);
  --itr1, --itr2;
  ans[D] += s[d][q] - s[d][old];
  ans[D] += c[q] * (d + 1 - *itr1) * (*next(itr1, 1) - d - 1);
  ans[D] -= c[old] * (d + 1 - *itr2) * (*next(itr2, 1) - d - 1);
  last[q].emplace(d + 1);
}

void greedy() {
  for (int i = 0; i < D - 2; ++i) {
    ans[i + 3] = INT_MIN;
    int maidx1 = 0, maidx2 = 0, maidx3 = 0;
    for (int k1 = 0; k1 < 26; ++k1) {
      int old1 = *last[k1].rbegin();
      last[k1].emplace(i + 1);
      for (int k2 = 0; k2 < 26; ++k2) {
        int old2 = *last[k2].rbegin();
        for (int k3 = 0; k3 < 26; ++k3) {
          int res = ans[i] + s[i][k1] + s[i + 1][k2] + s[i + 2][k3], res1, res2;
          int old3 = *last[k3].rbegin();
          for (int j = 0; j < 26; ++j)
            res -= c[j] * (i + 1 - *last[j].rbegin());
          res1 = res;
          last[k2].emplace(i + 2);
          for (int j = 0; j < 26; ++j)
            res -= c[j] * (i + 2 - *last[j].rbegin());
          res2 = res;
          last[k3].emplace(i + 3);
          for (int j = 0; j < 26; ++j)
            res -= c[j] * (i + 3 - *last[j].rbegin());

          if (ans[i + 3] < res) {
            ans[i + 3] = res;
            ans[i + 2] = res2 - s[i + 2][k3];
            ans[i + 1] = res1 - s[i + 2][k3] - s[i + 1][k2];
            maidx1 = k1;
            maidx2 = k2;
            maidx3 = k3;
          }
          last[k2].erase(i + 2);
          last[k3].erase(i + 3);
        }
      }
      last[k1].erase(i + 1);
    }
    t[i] = maidx1;
    last[t[i]].emplace(i + 1);
    if (i == D - 3) {
      t[i + 1] = maidx2;
      t[i + 2] = maidx3;
      last[t[i + 1]].emplace(i + 2);
      last[t[i + 2]].emplace(i + 3);
    }
  }
}

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);

  timer tm;

  cin >> D;
  for (int i = 0; i < 26; ++i)
    cin >> c[i];
  for (int i = 0; i < D; ++i)
    for (int j = 0; j < 26; ++j)
      cin >> s[i][j];

  for (int i = 0; i < 26; ++i)
    last[i].emplace(0);

  greedy();

  for (int i = 0; i < 26; ++i)
    last[i].emplace(D + 1);

  std::mt19937_64 gen(1);
  std::uniform_int_distribution<> dist1(1, D);
  std::uniform_int_distribution<> dist2(1, 26);

  while (tm.get() < 1900) {
    stack<pii> stk;
    int pres = ans[D];
    for (int j = 0; j < 20; ++j) {
      int d = dist1(gen), q = dist2(gen);
      d--, q--;
      stk.emplace(pii{d, t[d]});
      change(d, q);
      if (ans[D] >= pres)
        break;
    }
    if (ans[D] < pres) {
      while (!stk.empty()) {
        auto [d, q] = stk.top();
        change(d, q);
        stk.pop();
      }
    }
  }
  for (int i = 0; i < D; ++i)
    cout << t[i] + 1 << '\\n';

  return 0;
}