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

#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < (n); i++)
using namespace std;
using ll = long long;
using PL = pair<ll, ll>;
using P = pair<int, int>;
constexpr int INF = 1000000000;
constexpr long long HINF = 1000000000000000;
constexpr long long MOD = 1000000007; // = 998244353;
constexpr double EPS = 1e-4;
constexpr double PI = 3.14159265358979;

struct state {
  int D;
  vector<int> c;
  vector<vector<int>> s;
  vector<int> before;

  state(int D, vector<int> &c, vector<vector<int>> &s)
      : D(D), c(c), s(s), before(26, -1) {}

  vector<ll> solve() {
    ll ans = -INF, best_i = -1;
    vector<ll> indexs;
    rep(i, 26) {
      vector<ll> op = optimization(i);
      if (op[0] > ans) {
        ans = op[0];
        indexs = op;
      }
    }
    return indexs;
  }

  // あるパラメータの時の解
  vector<ll> optimization(int x) {
    vector<ll> indexs(D + 1);
    ll sco = 0;
    rep(i, D) {
      int ret, idx = next_day(i, x);
      before[idx] = i;
      indexs[i + 1] = idx + 1;
      sco += ret;
    }
    indexs[0] = sco;
    return indexs;
  }

  // d日目にi番目のコンテストを開いた時のx日後までのスコア
  int evaluate(int d, int i, int x) {
    int score = 0;
    score += s[d][i];
    int bf = before[i];
    before[i] = d;
    rep(k, x) rep(j, 26) { score -= (k + d - before[j]) * c[j]; }
    before[i] = bf;
    return score;
  }

  // d日目のスコア
  int next_day(int d, int x) {
    int idx = -1, ans = -INF;
    rep(i, 26) {
      // choose i th contest
      int ret = evaluate(d, i, x);
      if (ans < ret) {
        ans = ret;
        idx = i;
      }
      before[i] = x;
    }
    return ans, idx;
  }
};

int main() {
  int D;
  cin >> D;
  ;
  vector<int> c(26);
  rep(i, 26) cin >> c[i];
  vector<vector<int>> s(D, vector<int>(26));
  rep(i, D) rep(j, 26) cin >> s[i][j];

  /*
  //答え確認用
  printf("%d\\n",D);
  rep(i,26) printf("%d ",c[i]);
  cout << '\\n';
  rep(i,D){
      rep(j,26) printf("%d ",s[i][j]);
      cout << '\\n';
  }
  */

  state st(D, c, s);
  vector<ll> ans = st.solve();
  rep(i, D) printf("%ld\\n", ans[i + 1]);
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < (n); i++)
using namespace std;
using ll = long long;
using PL = pair<ll, ll>;
using P = pair<int, int>;
constexpr int INF = 1000000000;
constexpr long long HINF = 1000000000000000;
constexpr long long MOD = 1000000007; // = 998244353;
constexpr double EPS = 1e-4;
constexpr double PI = 3.14159265358979;

struct state {
  int D;
  vector<int> c;
  vector<vector<int>> s;
  vector<int> before;

  state(int D, vector<int> &c, vector<vector<int>> &s)
      : D(D), c(c), s(s), before(26, -1) {}

  vector<int> optimization() {
    vector<ll> ans(D + 1, 0);
    vector<int> indexs(D + 1);
    int i;
    for (i = 0; i < D - 4; i += 3) {
      vector<int> idxs = next_four_day(i);
      for (int j = 1; j <= 4; ++j) {
        before[idxs[j]] = i + j;
        indexs[i + j] = idxs[j] + 1;
      }
    }
    while (i < D) {
      int idx = next_day(i);
      before[idx] = i;
      indexs[i + 1] = idx + 1;
      i += 1;
    }
    return indexs;
  }

  int next_day(int d) {
    int idx = -1, ans = -INF;
    rep(i, 26) {
      int ret = s[d][i];
      int x = before[i];
      before[i] = d;
      rep(j, 26) { ret -= (d - before[j]) * c[j]; }
      if (ans < ret) {
        ans = ret;
        idx = i;
      }
      before[i] = x;
    }
    return idx;
  }

  vector<int> dfs(int i, int k, int bf, int now) {
    vector<int> rets(1);
    rets[0] = -INF;
    int idx = -1;

    // start
    if (k < 0) {
      rep(j, 26) {
        int x = before[j];
        before[j] = i;
        vector<int> ne = dfs(i + 1, j, bf, now);
        if (rets[0] < ne[0]) {
          idx = j;
          swap(rets, ne);
        }
        before[j] = x;
      }
      rets.push_back(idx);
      return rets;
    }

    // score
    bf += s[i][k];
    rep(j, 26) { bf -= (i - before[j]) * c[j]; }

    // finish
    if (now == i) {
      vector<int> rets(1);
      rets[0] = bf;
      return rets;
    }

    // nextday
    rep(j, 26) {
      int x = before[j];
      before[j] = i;
      vector<int> ne = dfs(i + 1, j, bf, now);
      if (rets[0] < ne[0]) {
        swap(rets, ne);
        idx = j;
      }
      before[j] = x;
    }
    rets.push_back(idx);
    return rets;
  }

  vector<int> next_four_day(int d) {
    vector<int> ans = dfs(d, -1, 0, d + 4);
    return ans;
  }
};

int main() {
  int D;
  cin >> D;
  ;
  vector<int> c(26);
  rep(i, 26) cin >> c[i];
  vector<vector<int>> s(D, vector<int>(26));
  rep(i, D) rep(j, 26) cin >> s[i][j];

  state st(D, c, s);
  vector<int> ans = st.optimization();
  rep(i, D) printf("%d\\n", ans[i + 1]);
  return 0;
}