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

#include <bits/stdc++.h>
using namespace std;
#define reps(i, s, n) for (int i = s; i < n; i++)
#define rep(i, n) reps(i, 0, n)
#define Rreps(i, n, e) for (int i = n - 1; i >= e; --i)
#define Rrep(i, n) Rreps(i, n, 0)
#define ALL(a) a.begin(), a.end()
#define fi first
#define se second
typedef long long ll;
typedef vector<ll> vec;
typedef vector<vec> mat;

ll N, M, H, W, Q, K, A, B;
string S;
const ll MOD = 998244353;
// const ll MOD = (1e+9) + 7;
typedef pair<ll, ll> P;
const ll INF = (1LL << 58);

template <class T> bool chmin(T &a, const T &b) {
  if (a > b) {
    a = b;
    return true;
  } else
    return false;
}
template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return true;
  } else
    return false;
}

ll calc(mat &s, vec &c, vec &ans) {
  ll res = 0;
  vec last(26, 0);
  rep(day, N) {
    last[ans[day]] = day + 1;
    rep(i, 26) {
      if (i == ans[day]) {
        res += s[day][i];
      } else {
        res -= c[i] * (day + 1 - last[i]);
      }
    }
  }
  return res;
}

void make_ans_another(vec &ans, mat &ans_another) {
  rep(day, N) { ans_another[ans[day]].push_back(day + 1); }
  rep(i, 26) ans_another[i].push_back(N + 1);
}

int main() {
  cin >> N;
  vec c(26), last(26, 0), ans(N);
  mat s(N, vec(26));
  rep(i, 26) cin >> c[i];
  rep(i, N) rep(j, 26) cin >> s[i][j];
  ll score = 0;
  rep(day, N) {
    vec temp(26, 0);
    rep(i, 26) temp[i] += s[day][i];
    rep(i, 26) temp[i] += c[i] * (day + 1 - last[i]);
    auto ite = max_element(ALL(temp));
    ans[day] = ite - temp.begin();
    last[ite - temp.begin()] = day + 1;
    rep(i, 26) {
      if (i == ans[day]) {
        score += s[day][i];
      } else {
        score -= c[i] * (day + 1 - last[i]);
      }
    }
  }

  ll total_num = 200000;
  bool stopped = false;
  while (total_num && stopped) {
    mat minus(N + 1, vec(26, 0)), ans_another(26, vec(1, 0));
    make_ans_another(ans, ans_another);
    ++total_num;
    priority_queue<P> pque;
    rep(i, 26) {
      vec &table = ans_another[i];
      int id = 0;
      reps(day, 1, N) {
        if (table[id + 1] == day) {
          ++id;
        } else {
          ll weight = s[day - 1][i], len = table[id + 1] - table[id] - 1;
          ll wh = day - table[id];
          weight +=
              c[i] * ((len * (len + 1) / 2) -
                      (wh * (wh - 1) / 2 + (len - wh) * (len - wh + 1) / 2));

          ll vs = ans[day - 1];
          auto ite = lower_bound(ALL(ans_another[vs]), day);
          len = *(++ite);
          ite -= 2;
          len -= *ite + 1;
          wh = day - *ite;
          weight -=
              s[day - 1][vs] +
              c[vs] * ((len * (len + 1) / 2) -
                       (wh * (wh - 1) / 2 + (len - wh) * (len - wh + 1) / 2));

          pque.emplace(weight, day * 26 + i);
        }
      }
    }
    if ((total_num -= 20) < 0)
      break;
    rep(check, 10000) {
      P p = pque.top();
      pque.pop();
      ll day = p.se / 26, i = p.se % 26;
      --day;
      if (p.se <= 0) {
        stopped = check == 0;
        break;
      } else if (check == 0) {
        ans[day] = i;
        score += p.se;
      } else {
        ll old = ans[day];
        ans[day] = i;
        if (!chmax(score, calc(s, c, ans))) {
          ans[day] = old;
          check *= 10;
        }
      }
      if ((--total_num) < 0)
        break;
    }
  }

  rep(i, N) cout << ans[i] + 1 << '\\n';
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#define reps(i, s, n) for (int i = s; i < n; i++)
#define rep(i, n) reps(i, 0, n)
#define Rreps(i, n, e) for (int i = n - 1; i >= e; --i)
#define Rrep(i, n) Rreps(i, n, 0)
#define ALL(a) a.begin(), a.end()
#define fi first
#define se second
typedef long long ll;
typedef vector<ll> vec;
typedef vector<vec> mat;

ll N, M, H, W, Q, K, A, B;
string S;
const ll MOD = 998244353;
// const ll MOD = (1e+9) + 7;
typedef pair<ll, ll> P;
const ll INF = (1LL << 58);

template <class T> bool chmin(T &a, const T &b) {
  if (a > b) {
    a = b;
    return true;
  } else
    return false;
}
template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return true;
  } else
    return false;
}

ll calc(mat &s, vec &c, vec &ans) {
  ll res = 0;
  vec last(26, 0);
  rep(day, N) {
    last[ans[day]] = day + 1;
    rep(i, 26) {
      if (i == ans[day]) {
        res += s[day][i];
      } else {
        res -= c[i] * (day + 1 - last[i]);
      }
    }
  }
  return res;
}

void make_ans_another(vec &ans, mat &ans_another) {
  rep(day, N) { ans_another[ans[day]].push_back(day + 1); }
  rep(i, 26) ans_another[i].push_back(N + 1);
}

int main() {
  cin >> N;
  vec c(26), last(26, 0), ans(N);
  mat s(N, vec(26));
  rep(i, 26) cin >> c[i];
  rep(i, N) rep(j, 26) cin >> s[i][j];
  ll score = 0;
  rep(day, N) {
    vec temp(26, 0);
    rep(i, 26) temp[i] += s[day][i];
    rep(i, 26) temp[i] += c[i] * (day + 1 - last[i]);
    auto ite = max_element(ALL(temp));
    ans[day] = ite - temp.begin();
    last[ite - temp.begin()] = day + 1;
    rep(i, 26) {
      if (i == ans[day]) {
        score += s[day][i];
      } else {
        score -= c[i] * (day + 1 - last[i]);
      }
    }
  }

  ll total_num = 200000;
  bool stopped = false;
  while (total_num || stopped) {
    mat minus(N + 1, vec(26, 0)), ans_another(26, vec(1, 0));
    make_ans_another(ans, ans_another);
    ++total_num;
    priority_queue<P> pque;
    rep(i, 26) {
      vec &table = ans_another[i];
      int id = 0;
      reps(day, 1, N) {
        if (table[id + 1] == day) {
          ++id;
        } else {
          ll weight = s[day - 1][i], len = table[id + 1] - table[id] - 1;
          ll wh = day - table[id];
          weight +=
              c[i] * ((len * (len + 1) / 2) -
                      (wh * (wh - 1) / 2 + (len - wh) * (len - wh + 1) / 2));

          ll vs = ans[day - 1];
          auto ite = lower_bound(ALL(ans_another[vs]), day);
          len = *(++ite);
          ite -= 2;
          len -= *ite + 1;
          wh = day - *ite;
          weight -=
              s[day - 1][vs] +
              c[vs] * ((len * (len + 1) / 2) -
                       (wh * (wh - 1) / 2 + (len - wh) * (len - wh + 1) / 2));

          pque.emplace(weight, day * 26 + i);
        }
      }
    }
    if ((total_num -= 20) < 0)
      break;
    rep(check, 10000) {
      P p = pque.top();
      pque.pop();
      ll day = p.se / 26, i = p.se % 26;
      --day;
      if (p.se <= 0) {
        stopped = check == 0;
        break;
      } else if (check == 0) {
        ans[day] = i;
        score += p.se;
      } else {
        ll old = ans[day];
        ans[day] = i;
        if (!chmax(score, calc(s, c, ans))) {
          ans[day] = old;
          check *= 10;
        }
      }
      if ((--total_num) < 0)
        break;
    }
  }

  rep(i, N) cout << ans[i] + 1 << '\\n';
}