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

// #define DBG2

#include <bits/stdc++.h>
#include <random>

using namespace std;
const int N_CNTST = 26;
int D;
vector<long long> C;
vector<vector<long long>> S;
long long CSum = 0;

vector<int> LAST(26, 0);

vector<int> T;
int main() {
  cin >> D;
  cin.ignore();

  long long tmp;
  for (int i = 0; i < N_CNTST; i++) {
    cin >> tmp;
    C.push_back(tmp);
    CSum += tmp;
  }
  cin.ignore();
  for (int i = 0; i < D; i++) {
    vector<long long> S_;
    for (int j = 0; j < N_CNTST; j++) {
      cin >> tmp;
      S_.push_back(tmp);
    }
    cin.ignore();
    S.push_back(S_);
  }

  // tを算出
  long long sum = 0;
  long long msum = 0; // 減算Sum
  for (int i = 0; i < D; i++) {
    // cin >> tmp;
    // T.push_back((int)tmp);
    // cin.ignore();

    int idx;
    long long ssum;
    long long smsum;
    long long tmp_msum;
    long long max_sum = LLONG_MIN;
    for (int t_idx = 0; t_idx < N_CNTST; t_idx++) {
      smsum = msum + CSum - C[t_idx] * (i - LAST[t_idx] + 1);
      ssum = sum + S[i][t_idx] - smsum;
      if (max_sum < ssum || t_idx == 0) { // どれもいっしょかつマイナスの場合
        idx = t_idx;
        max_sum = ssum;
        tmp_msum = smsum;
      }
    }
    msum = tmp_msum;
    LAST[idx] = i + 1;
    sum = max_sum;
    // cout << idx + 1 << endl; // 0-origin -> 1-origin
    T.push_back(idx + 1);
    // cout << sum << endl;
  }

#ifdef DBG
  printf("T is:");
  for (int i = 0; i < T.size(); i++) {
    printf("%d:%d, ", i, T[i]);
  }
  printf("\\n");
#endif

  int M;
#ifdef DBG2
  cin >> M;
#else
  M = 1000000;
#endif
  cin.ignore();
  std::random_device rnd; // 非決定的な乱数生成器
  for (int i = 0; i < M; i++) {

    // o_idxと、i_idxをランダムに決める
    unsigned long long rnd1 = rnd();
    unsigned long long rnd2 = rnd();
    int idx1 = rnd1 / (double)UINT32_MAX * (D - 1);
    int idx2 = rnd2 / (double)UINT32_MAX * (D - 1);
#ifdef DBG2
    cout << "idx1 " << idx1 << " idx2 " << idx2 << endl;
#endif
    if (idx1 == idx2)
      continue;
    if (T[idx1] == T[idx2])
      continue;
    vector<int> T_ = T;
    long long sum_ = sum;
    {
      int o_idx = idx1;
      int chg_idx = idx2;
      int chg = T_[chg_idx];

      int prb = T_[o_idx];
      int o_s_idx;
      bool s_found = false;
      for (int j = o_idx - 1; j >= 0; j--) {
        if (prb == T_[j]) {
          o_s_idx = j;
          s_found = true;
          break;
        }
      }
      if (!s_found)
        o_s_idx = -1;
      int o_e_idx;
      bool e_found = false;
      for (int j = o_idx + 1; j < D; j++) {
        if (prb == T_[j]) {
          o_e_idx = j;
          e_found = true;
          break;
        }
      }
      if (!e_found)
        o_e_idx = D;
      int chg_s_idx;
      bool chg_s_found = false;
      for (int j = o_idx - 1; j >= 0; j--) {
        if (chg == T_[j]) {
          chg_s_idx = j;
          chg_s_found = true;
          break;
        }
      }
      if (!chg_s_found)
        chg_s_idx = -1;
      int chg_e_idx;
      bool chg_e_found = false;
      for (int j = o_idx + 1; j < D; j++) {
        if (chg == T_[j]) {
          chg_e_idx = j;
          chg_e_found = true;
          break;
        }
      }
      if (!chg_e_found)
        chg_e_idx = D;

      int s_days = o_idx - o_s_idx;
      // int s_n_sum = s_days * (1 + s_days) / 2;
      int s_n_sum = s_days;

#ifdef DBG
      printf("s_days %d o_idx %d o_s_idx %d\\n", s_days, o_idx, o_s_idx);
#endif
      int e_days = o_e_idx - o_idx;
#ifdef DBG
      printf("e_days %d o_e_idx %d o_idx %d\\n", e_days, o_e_idx, o_idx);
#endif
      sum_ -= s_n_sum * C[T_[o_idx] - 1] * e_days;
#ifdef DBG
      printf("tempsum %lld s_n_sum %d COST %lld e_days %d\\n", sum, s_n_sum,
             C[T[o_idx] - 1], e_days);
#endif

      int chg_s_days = o_idx - chg_s_idx;
#ifdef DBG
      printf("chg_s_days %d o_idx %d chg_s_idx %d\\n", chg_s_days, o_idx,
             chg_s_idx);
#endif
      // int chg_n_sum = chg_s_days * (1 + chg_s_days) / 2;
      int chg_n_sum = chg_s_days;
      int chg_e_days = chg_e_idx - o_idx;
#ifdef DBG
      printf("chg_e_days %d chg_e_idx %d o_idx %d\\n", chg_e_days, chg_e_idx,
             o_idx);
#endif
      sum_ += chg_n_sum * C[chg - 1] * chg_e_days;
#ifdef DBG
      printf("tempsum %lld chg_n_sum %d COST %lld chg_e_days %d\\n", sum,
             chg_n_sum, C[chg - 1], chg_e_days);
#endif
      sum_ += -S[o_idx][T_[o_idx] - 1] + S[o_idx][chg - 1];
#ifdef DBG
      cout << sum << endl;
#endif
      T_[o_idx] = chg;
    }

    {
      int o_idx = idx2;
      int chg_idx = idx1;
      int chg = T_[chg_idx];

      int prb = T_[o_idx];
      int o_s_idx;
      bool s_found = false;
      for (int j = o_idx - 1; j >= 0; j--) {
        if (prb == T_[j]) {
          o_s_idx = j;
          s_found = true;
          break;
        }
      }
      if (!s_found)
        o_s_idx = -1;
      int o_e_idx;
      bool e_found = false;
      for (int j = o_idx + 1; j < D; j++) {
        if (prb == T_[j]) {
          o_e_idx = j;
          e_found = true;
          break;
        }
      }
      if (!e_found)
        o_e_idx = D;
      int chg_s_idx;
      bool chg_s_found = false;
      for (int j = o_idx - 1; j >= 0; j--) {
        if (chg == T_[j]) {
          chg_s_idx = j;
          chg_s_found = true;
          break;
        }
      }
      if (!chg_s_found)
        chg_s_idx = -1;
      int chg_e_idx;
      bool chg_e_found = false;
      for (int j = o_idx + 1; j < D; j++) {
        if (chg == T_[j]) {
          chg_e_idx = j;
          chg_e_found = true;
          break;
        }
      }
      if (!chg_e_found)
        chg_e_idx = D;

      int s_days = o_idx - o_s_idx;
      // int s_n_sum = s_days * (1 + s_days) / 2;
      int s_n_sum = s_days;

#ifdef DBG
      printf("s_days %d o_idx %d o_s_idx %d\\n", s_days, o_idx, o_s_idx);
#endif
      int e_days = o_e_idx - o_idx;
#ifdef DBG
      printf("e_days %d o_e_idx %d o_idx %d\\n", e_days, o_e_idx, o_idx);
#endif
      sum_ -= s_n_sum * C[T_[o_idx] - 1] * e_days;
#ifdef DBG
      printf("tempsum %lld s_n_sum %d COST %lld e_days %d\\n", sum, s_n_sum,
             C[T[o_idx] - 1], e_days);
#endif

      int chg_s_days = o_idx - chg_s_idx;
#ifdef DBG
      printf("chg_s_days %d o_idx %d chg_s_idx %d\\n", chg_s_days, o_idx,
             chg_s_idx);
#endif
      // int chg_n_sum = chg_s_days * (1 + chg_s_days) / 2;
      int chg_n_sum = chg_s_days;
      int chg_e_days = chg_e_idx - o_idx;
#ifdef DBG
      printf("chg_e_days %d chg_e_idx %d o_idx %d\\n", chg_e_days, chg_e_idx,
             o_idx);
#endif
      sum_ += chg_n_sum * C[chg - 1] * chg_e_days;
#ifdef DBG
      printf("tempsum %lld chg_n_sum %d COST %lld chg_e_days %d\\n", sum,
             chg_n_sum, C[chg - 1], chg_e_days);
#endif
      sum_ += -S[o_idx][T_[o_idx] - 1] + S[o_idx][chg - 1];
#ifdef DBG
      cout << sum << endl;
#endif
      T_[o_idx] = chg;
    }

    if (sum < sum_) {
      sum = sum_;
      T = T_;
    }
  }
  for (int i = 0; i < T.size(); i++) {
    cout << T[i] << endl;
  }
  return 0;
}


Error: TLE

Fixed submission:

#define DBG2

#include <bits/stdc++.h>
#include <random>

using namespace std;
const int N_CNTST = 26;
int D;
vector<long long> C;
vector<vector<long long>> S;
long long CSum = 0;

vector<int> LAST(26, 0);

vector<int> T;
int main() {
  cin >> D;
  cin.ignore();

  long long tmp;
  for (int i = 0; i < N_CNTST; i++) {
    cin >> tmp;
    C.push_back(tmp);
    CSum += tmp;
  }
  cin.ignore();
  for (int i = 0; i < D; i++) {
    vector<long long> S_;
    for (int j = 0; j < N_CNTST; j++) {
      cin >> tmp;
      S_.push_back(tmp);
    }
    cin.ignore();
    S.push_back(S_);
  }

  // tを算出
  long long sum = 0;
  long long msum = 0; // 減算Sum
  for (int i = 0; i < D; i++) {
    // cin >> tmp;
    // T.push_back((int)tmp);
    // cin.ignore();

    int idx;
    long long ssum;
    long long smsum;
    long long tmp_msum;
    long long max_sum = LLONG_MIN;
    for (int t_idx = 0; t_idx < N_CNTST; t_idx++) {
      smsum = msum + CSum - C[t_idx] * (i - LAST[t_idx] + 1);
      ssum = sum + S[i][t_idx] - smsum;
      if (max_sum < ssum || t_idx == 0) { // どれもいっしょかつマイナスの場合
        idx = t_idx;
        max_sum = ssum;
        tmp_msum = smsum;
      }
    }
    msum = tmp_msum;
    LAST[idx] = i + 1;
    sum = max_sum;
    // cout << idx + 1 << endl; // 0-origin -> 1-origin
    T.push_back(idx + 1);
    // cout << sum << endl;
  }

#ifdef DBG
  printf("T is:");
  for (int i = 0; i < T.size(); i++) {
    printf("%d:%d, ", i, T[i]);
  }
  printf("\\n");
#endif

  int M;
  cin >> M;
  cin.ignore();
  std::random_device rnd; // 非決定的な乱数生成器
  for (int i = 0; i < M; i++) {

    // o_idxと、i_idxをランダムに決める
    unsigned long long rnd1 = rnd();
    unsigned long long rnd2 = rnd();
    int idx1 = rnd1 / (double)UINT32_MAX * (D - 1);
    int idx2 = rnd2 / (double)UINT32_MAX * (D - 1);
#ifdef DBG2
    cout << "idx1 " << idx1 << " idx2 " << idx2 << endl;
#endif
    if (idx1 == idx2)
      continue;
    if (T[idx1] == T[idx2])
      continue;
    vector<int> T_ = T;
    long long sum_ = sum;
    {
      int o_idx = idx1;
      int chg_idx = idx2;
      int chg = T_[chg_idx];

      int prb = T_[o_idx];
      int o_s_idx;
      bool s_found = false;
      for (int j = o_idx - 1; j >= 0; j--) {
        if (prb == T_[j]) {
          o_s_idx = j;
          s_found = true;
          break;
        }
      }
      if (!s_found)
        o_s_idx = -1;
      int o_e_idx;
      bool e_found = false;
      for (int j = o_idx + 1; j < D; j++) {
        if (prb == T_[j]) {
          o_e_idx = j;
          e_found = true;
          break;
        }
      }
      if (!e_found)
        o_e_idx = D;
      int chg_s_idx;
      bool chg_s_found = false;
      for (int j = o_idx - 1; j >= 0; j--) {
        if (chg == T_[j]) {
          chg_s_idx = j;
          chg_s_found = true;
          break;
        }
      }
      if (!chg_s_found)
        chg_s_idx = -1;
      int chg_e_idx;
      bool chg_e_found = false;
      for (int j = o_idx + 1; j < D; j++) {
        if (chg == T_[j]) {
          chg_e_idx = j;
          chg_e_found = true;
          break;
        }
      }
      if (!chg_e_found)
        chg_e_idx = D;

      int s_days = o_idx - o_s_idx;
      // int s_n_sum = s_days * (1 + s_days) / 2;
      int s_n_sum = s_days;

#ifdef DBG
      printf("s_days %d o_idx %d o_s_idx %d\\n", s_days, o_idx, o_s_idx);
#endif
      int e_days = o_e_idx - o_idx;
#ifdef DBG
      printf("e_days %d o_e_idx %d o_idx %d\\n", e_days, o_e_idx, o_idx);
#endif
      sum_ -= s_n_sum * C[T_[o_idx] - 1] * e_days;
#ifdef DBG
      printf("tempsum %lld s_n_sum %d COST %lld e_days %d\\n", sum, s_n_sum,
             C[T[o_idx] - 1], e_days);
#endif

      int chg_s_days = o_idx - chg_s_idx;
#ifdef DBG
      printf("chg_s_days %d o_idx %d chg_s_idx %d\\n", chg_s_days, o_idx,
             chg_s_idx);
#endif
      // int chg_n_sum = chg_s_days * (1 + chg_s_days) / 2;
      int chg_n_sum = chg_s_days;
      int chg_e_days = chg_e_idx - o_idx;
#ifdef DBG
      printf("chg_e_days %d chg_e_idx %d o_idx %d\\n", chg_e_days, chg_e_idx,
             o_idx);
#endif
      sum_ += chg_n_sum * C[chg - 1] * chg_e_days;
#ifdef DBG
      printf("tempsum %lld chg_n_sum %d COST %lld chg_e_days %d\\n", sum,
             chg_n_sum, C[chg - 1], chg_e_days);
#endif
      sum_ += -S[o_idx][T_[o_idx] - 1] + S[o_idx][chg - 1];
#ifdef DBG
      cout << sum << endl;
#endif
      T_[o_idx] = chg;
    }

    {
      int o_idx = idx2;
      int chg_idx = idx1;
      int chg = T_[chg_idx];

      int prb = T_[o_idx];
      int o_s_idx;
      bool s_found = false;
      for (int j = o_idx - 1; j >= 0; j--) {
        if (prb == T_[j]) {
          o_s_idx = j;
          s_found = true;
          break;
        }
      }
      if (!s_found)
        o_s_idx = -1;
      int o_e_idx;
      bool e_found = false;
      for (int j = o_idx + 1; j < D; j++) {
        if (prb == T_[j]) {
          o_e_idx = j;
          e_found = true;
          break;
        }
      }
      if (!e_found)
        o_e_idx = D;
      int chg_s_idx;
      bool chg_s_found = false;
      for (int j = o_idx - 1; j >= 0; j--) {
        if (chg == T_[j]) {
          chg_s_idx = j;
          chg_s_found = true;
          break;
        }
      }
      if (!chg_s_found)
        chg_s_idx = -1;
      int chg_e_idx;
      bool chg_e_found = false;
      for (int j = o_idx + 1; j < D; j++) {
        if (chg == T_[j]) {
          chg_e_idx = j;
          chg_e_found = true;
          break;
        }
      }
      if (!chg_e_found)
        chg_e_idx = D;

      int s_days = o_idx - o_s_idx;
      // int s_n_sum = s_days * (1 + s_days) / 2;
      int s_n_sum = s_days;

#ifdef DBG
      printf("s_days %d o_idx %d o_s_idx %d\\n", s_days, o_idx, o_s_idx);
#endif
      int e_days = o_e_idx - o_idx;
#ifdef DBG
      printf("e_days %d o_e_idx %d o_idx %d\\n", e_days, o_e_idx, o_idx);
#endif
      sum_ -= s_n_sum * C[T_[o_idx] - 1] * e_days;
#ifdef DBG
      printf("tempsum %lld s_n_sum %d COST %lld e_days %d\\n", sum, s_n_sum,
             C[T[o_idx] - 1], e_days);
#endif

      int chg_s_days = o_idx - chg_s_idx;
#ifdef DBG
      printf("chg_s_days %d o_idx %d chg_s_idx %d\\n", chg_s_days, o_idx,
             chg_s_idx);
#endif
      // int chg_n_sum = chg_s_days * (1 + chg_s_days) / 2;
      int chg_n_sum = chg_s_days;
      int chg_e_days = chg_e_idx - o_idx;
#ifdef DBG
      printf("chg_e_days %d chg_e_idx %d o_idx %d\\n", chg_e_days, chg_e_idx,
             o_idx);
#endif
      sum_ += chg_n_sum * C[chg - 1] * chg_e_days;
#ifdef DBG
      printf("tempsum %lld chg_n_sum %d COST %lld chg_e_days %d\\n", sum,
             chg_n_sum, C[chg - 1], chg_e_days);
#endif
      sum_ += -S[o_idx][T_[o_idx] - 1] + S[o_idx][chg - 1];
#ifdef DBG
      cout << sum << endl;
#endif
      T_[o_idx] = chg;
    }

    if (sum < sum_) {
      sum = sum_;
      T = T_;
    }
  }
  for (int i = 0; i < T.size(); i++) {
    cout << T[i] << endl;
  }
  return 0;
}