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

#pragma GCC optimize("Ofast")
#include <bits/stdc++.h>

#define ll long long
#define ull unsigned long long

using namespace std;

int main() {
  ll D;
  cin >> D;

  ll N = 26;
  vector<ll> c(N);
  for (ll ii = 0; ii < N; ++ii) {
    cin >> c[ii];
  }
  vector<vector<ll>> s(D, vector<ll>(N));
  for (ll ii = 0; ii < D; ++ii) {
    for (ll jj = 0; jj < N; ++jj) {
      cin >> s[ii][jj];
    }
  }

  // 貪欲法
  vector<int> ans(D, 0);
  vector<int> last(N, -1);
  vector<set<int>> select(N);
  for (int ii = 0; ii < N; ++ii) {
    select[ii].insert(-1);
  }
  for (ll d = 0; d < D; ++d) {
    // d日目の最適解を選ぶ
    ll val = 0;
    ll idx = 0;
    for (int ii = 0; ii < N; ++ii) {
      ll tmp = s[d][ii] + c[ii] * (d - last[ii]);
      if (tmp > val) {
        val = tmp;
        idx = ii;
      }
    }
    ans[d] = idx;
    last[idx] = d;
    select[idx].insert(d);
  }

  // ランダムに値を変更する
  int trial = 1001001;
  for (int t = 0; t < trial; ++t) {
    int d = rand() % D;
    int prevType = ans[d];
    int newType = rand() % N;
    while (prevType == newType) {
      newType = rand() % N;
    }
    ll dif = s[d][newType] - s[d][prevType];

    auto itr = upper_bound(select[newType].begin(), select[newType].end(), d);
    ll newDist = (itr != select[newType].end()) ? (*itr - d) : (D - d);
    auto itr2 =
        upper_bound(select[newType].rbegin(), select[newType].rend(), d);
    ll newVal = d - *itr2;
    dif += newVal * newDist;

    itr = upper_bound(select[prevType].begin(), select[prevType].end(), d);
    ll prevDist = (itr != select[prevType].end()) ? (*itr - d) : (D - d);
    itr2 = upper_bound(select[prevType].rbegin(), select[prevType].rend(), d);
    ll prevVal = d - *itr2;
    dif -= prevVal * prevDist;

    if (dif > 0) {
      select[prevType].erase(d);
      select[newType].insert(d);
      ans[d] = newType;
    }
  }

  for (ll d = 0; d < D; ++d) {
    cout << ans[d] + 1 << "\\n";
  }

  return 0;
}


Error: TLE

Fixed submission:

#pragma GCC optimize("Ofast")
#include <bits/stdc++.h>

#define ll long long
#define ull unsigned long long

using namespace std;

int main() {
  ll D;
  cin >> D;

  ll N = 26;
  vector<ll> c(N);
  for (ll ii = 0; ii < N; ++ii) {
    cin >> c[ii];
  }
  vector<vector<ll>> s(D, vector<ll>(N));
  for (ll ii = 0; ii < D; ++ii) {
    for (ll jj = 0; jj < N; ++jj) {
      cin >> s[ii][jj];
    }
  }

  // 貪欲法
  vector<int> ans(D, 0);
  vector<int> last(N, -1);
  vector<set<int>> select(N);
  for (int ii = 0; ii < N; ++ii) {
    select[ii].insert(-1);
  }
  for (ll d = 0; d < D; ++d) {
    // d日目の最適解を選ぶ
    ll val = 0;
    ll idx = 0;
    for (int ii = 0; ii < N; ++ii) {
      ll tmp = s[d][ii] + c[ii] * (d - last[ii]);
      if (tmp > val) {
        val = tmp;
        idx = ii;
      }
    }
    ans[d] = idx;
    last[idx] = d;
    select[idx].insert(d);
  }

  // ランダムに値を変更する
  int trial = 2001001;
  for (int t = 0; t < trial; ++t) {
    int d = rand() % D;
    int prevType = ans[d];
    int newType = rand() % N;
    while (prevType == newType) {
      newType = rand() % N;
    }
    ll dif = s[d][newType] - s[d][prevType];

    auto itr = upper_bound(select[newType].begin(), select[newType].end(), d);
    ll newDist = (itr != select[newType].end()) ? (*itr - d) : (D - d);
    auto itr2 =
        upper_bound(select[newType].rbegin(), select[newType].rend(), d);
    ll newVal = d - *itr2;
    dif += newVal * newDist;

    itr = upper_bound(select[prevType].begin(), select[prevType].end(), d);
    ll prevDist = (itr != select[prevType].end()) ? (*itr - d) : (D - d);
    itr2 = upper_bound(select[prevType].rbegin(), select[prevType].rend(), d);
    ll prevVal = d - *itr2;
    dif -= prevVal * prevDist;

    if (dif > 0) {
      select[prevType].erase(d);
      select[newType].insert(d);
      ans[d] = newType;
    }
  }

  for (ll d = 0; d < D; ++d) {
    cout << ans[d] + 1 << "\\n";
  }

  return 0;
}