Problem p02618 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>
#define int long long
#define rep(i, n) for (int i = 0; i < (int)(n); i++)
using namespace std;
constexpr int MOD = 1000000007;
constexpr int INF = numeric_limits<int>::max() / 2;
typedef pair<int, int> P;
using Graph = vector<vector<int>>;
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}
const long double PI = acos(-1.0);
// cout << fixed << setprecision(14);

set<int> st[26];

signed main() {
  cin.tie(nullptr);
  ios::sync_with_stdio(false);

  int D;
  cin >> D;
  rep(i, 26) {
    st[i].insert(-1);
    st[i].insert(D);
  }
  int c[26];
  int s[D][26];
  int t[D];
  int ok[26];
  rep(i, 26) ok[i] = -1;
  rep(i, 26) cin >> c[i];
  rep(i, D) {
    rep(j, 26) { cin >> s[i][j]; }
  }
  int num;
  rep(i, D) {
    int maxi = -INF;
    rep(k, 26) {
      int ans = 0;
      ans += s[i][k];
      // ok[t[i]] = i;
      rep(j, 26) {
        if (j == k)
          continue;
        ans -= c[j] * (i - ok[j]);
      }
      if (maxi < ans) {
        maxi = ans;
        num = k;
      }
    }
    ok[num] = i;
    t[i] = num;
    st[num].insert(i);
  }
  rep(loop, 3) {
    rep(i, D - 1) {
      int dx = 0;
      int old1 = t[i], old2 = t[i + 1];
      // st[_].erase(d[i]);
      auto itr = st[old1].upper_bound(i);
      int W2 = *itr;
      --itr;
      --itr;
      int W1 = *itr;
      dx -= c[old1] * (i - W1) * (W2 - i);
      auto itr1 = st[old2].upper_bound(i);
      int A2 = *itr1;
      --itr1;
      int A1 = *itr1;
      dx += c[old2] * (i - A1) * (A2 - i);
      // st[q[i]].insert(d[i]);
      dx += s[i][t[i + 1]] - s[i][t[i]];

      ++i;
      // st[_].erase(d[i]);
      itr = st[old2].upper_bound(i);
      W2 = *itr;
      --itr;
      --itr;
      W1 = *itr;
      dx -= c[old2] * (i - W1) * (W2 - i);
      itr1 = st[old1].upper_bound(i);
      A2 = *itr1;
      --itr1;
      A1 = *itr1;
      dx += c[old1] * (i - A1) * (A2 - i);
      // st[q[i]].insert(d[i]);
      dx += s[i][t[i - 1]] - s[i][t[i]];
      --i;
      if (dx > 0) {
        t[i] = old2;
        t[i + 1] = old1;
      }
    }
  }
  rep(loop, 100000) {
    rep(i, D - 1) {
      int dx = 0;
      int old1 = t[i], old2 = t[i + 1];
      // st[_].erase(d[i]);
      auto itr = st[old1].upper_bound(i);
      int W2 = *itr;
      --itr;
      --itr;
      int W1 = *itr;
      dx -= c[old1] * (i - W1) * (W2 - i);
      auto itr1 = st[old2].upper_bound(i);
      int A2 = *itr1;
      --itr1;
      int A1 = *itr1;
      dx += c[old2] * (i - A1) * (A2 - i);
      // st[q[i]].insert(d[i]);
      dx += s[i][t[i + 1]] - s[i][t[i]];

      ++i;
      // st[_].erase(d[i]);
      itr = st[old2].upper_bound(i);
      W2 = *itr;
      --itr;
      --itr;
      W1 = *itr;
      dx -= c[old2] * (i - W1) * (W2 - i);
      itr1 = st[old1].upper_bound(i);
      A2 = *itr1;
      --itr1;
      A1 = *itr1;
      dx += c[old1] * (i - A1) * (A2 - i);
      // st[q[i]].insert(d[i]);
      dx += s[i][t[i - 1]] - s[i][t[i]];
      --i;
      if (dx > 0) {
        t[i] = old2;
        t[i + 1] = old1;
      }
    }
  }
  rep(i, D) cout << t[i] + 1 << endl;
}

Fixed submission:
#include <bits/stdc++.h>
#define int long long
#define rep(i, n) for (int i = 0; i < (int)(n); i++)
using namespace std;
constexpr int MOD = 1000000007;
constexpr int INF = numeric_limits<int>::max() / 2;
typedef pair<int, int> P;
using Graph = vector<vector<int>>;
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}
const long double PI = acos(-1.0);
// cout << fixed << setprecision(14);

set<int> st[26];

signed main() {
  cin.tie(nullptr);
  ios::sync_with_stdio(false);

  int D;
  cin >> D;
  rep(i, 26) {
    st[i].insert(-1);
    st[i].insert(D);
  }
  int c[26];
  int s[D][26];
  int t[D];
  int ok[26];
  rep(i, 26) ok[i] = -1;
  rep(i, 26) cin >> c[i];
  rep(i, D) {
    rep(j, 26) { cin >> s[i][j]; }
  }
  int num;
  rep(i, D) {
    int maxi = -INF;
    rep(k, 26) {
      int ans = 0;
      ans += s[i][k];
      // ok[t[i]] = i;
      rep(j, 26) {
        if (j == k)
          continue;
        ans -= c[j] * (i - ok[j]);
      }
      if (maxi < ans) {
        maxi = ans;
        num = k;
      }
    }
    ok[num] = i;
    t[i] = num;
    st[num].insert(i);
  }
  rep(loop, 3) {
    rep(i, D - 1) {
      int dx = 0;
      int old1 = t[i], old2 = t[i + 1];
      // st[_].erase(d[i]);
      auto itr = st[old1].upper_bound(i);
      int W2 = *itr;
      --itr;
      --itr;
      int W1 = *itr;
      dx -= c[old1] * (i - W1) * (W2 - i);
      auto itr1 = st[old2].upper_bound(i);
      int A2 = *itr1;
      --itr1;
      int A1 = *itr1;
      dx += c[old2] * (i - A1) * (A2 - i);
      // st[q[i]].insert(d[i]);
      dx += s[i][t[i + 1]] - s[i][t[i]];

      ++i;
      // st[_].erase(d[i]);
      itr = st[old2].upper_bound(i);
      W2 = *itr;
      --itr;
      --itr;
      W1 = *itr;
      dx -= c[old2] * (i - W1) * (W2 - i);
      itr1 = st[old1].upper_bound(i);
      A2 = *itr1;
      --itr1;
      A1 = *itr1;
      dx += c[old1] * (i - A1) * (A2 - i);
      // st[q[i]].insert(d[i]);
      dx += s[i][t[i - 1]] - s[i][t[i]];
      --i;
      if (dx > 0) {
        t[i] = old2;
        t[i + 1] = old1;
      }
    }
  }
  rep(loop, 80000) {
    rep(i, D - 1) {
      int dx = 0;
      int old1 = t[i], old2 = t[i + 1];
      // st[_].erase(d[i]);
      auto itr = st[old1].upper_bound(i);
      int W2 = *itr;
      --itr;
      --itr;
      int W1 = *itr;
      dx -= c[old1] * (i - W1) * (W2 - i);
      auto itr1 = st[old2].upper_bound(i);
      int A2 = *itr1;
      --itr1;
      int A1 = *itr1;
      dx += c[old2] * (i - A1) * (A2 - i);
      // st[q[i]].insert(d[i]);
      dx += s[i][t[i + 1]] - s[i][t[i]];

      ++i;
      // st[_].erase(d[i]);
      itr = st[old2].upper_bound(i);
      W2 = *itr;
      --itr;
      --itr;
      W1 = *itr;
      dx -= c[old2] * (i - W1) * (W2 - i);
      itr1 = st[old1].upper_bound(i);
      A2 = *itr1;
      --itr1;
      A1 = *itr1;
      dx += c[old1] * (i - A1) * (A2 - i);
      // st[q[i]].insert(d[i]);
      dx += s[i][t[i - 1]] - s[i][t[i]];
      --i;
      if (dx > 0) {
        t[i] = old2;
        t[i + 1] = old1;
      }
    }
  }
  rep(i, D) cout << t[i] + 1 << endl;
}