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

#include <algorithm>
#include <bitset>
#include <cmath>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <string>
#include <vector>
#define vll vector<ll>
#define vvvl vector<vvl>
#define vvl vector<vector<ll>>
#define VV(a, b, c, d) vector<vector<d>>(a, vector<d>(b, c))
#define VVV(a, b, c, d) vector<vvl>(a, vvl(b, vll(c, d)));
#define re(c, b) for (ll c = 0; c < b; c++)
#define all(obj) (obj).begin(), (obj).end()
typedef long long int ll;
typedef long double ld;
using namespace std;

#include <random>
std::random_device seed_gen;
std::mt19937_64 engine(seed_gen());
// ull a = engine();
ll D;
vll now, c(26);
vvl s;
vector<set<ll>> st(26);
ll now_score = 0;

void init(ll d) {
  D = d;
  now.resize(d, 0);
  s = VV(d, 26, 0, ll);
}
void make_set(vll t) {
  for (int i = 0; i < D; i++)
    st[t[i]].insert(i);
  for (int i = 0; i < 26; i++)
    st[i].insert(-1);
  for (int i = 0; i < 26; i++)
    st[i].insert(D);
}
ll calc(vll t) {
  ll score = 0;
  vll cnt(26, 0);
  for (int i = 0; i < D; i++) {
    ll idx = t[i];
    for (int j = 0; j < 26; j++) {
      if (j == idx)
        cnt[j] = 0;
      else {
        cnt[j]++;
        score -= cnt[j] * c[j];
      }
    }
    score += s[i][idx];
  }
  return score;
}

ll ch_score(ll day, ll after) {
  ll score = now_score;
  ll prev = now[day]; // 前に選択していたもの
  if (prev == after)
    return score;

  // prevを消したことで減るスコア
  score -= s[day][prev];

  auto itr = st[prev].find(day);
  auto p = itr, q = itr;
  ll x = day - (*--p) - 1;
  ll y = (*++q) - day - 1;

  ll z = x + y + 1;
  // x, y分増え、z分減る
  score += c[prev] * ((x * (x + 1)) / 2 + (y * (y + 1)) / 2);
  score -= c[prev] * ((z * (z + 1)) / 2);
  // afterを挿入した事で増えるスコア
  score += s[day][after];
  itr = st[after].upper_bound(day);
  p = itr;
  p--;

  z = (*itr) - (*p) - 1;
  x = day - (*p) - 1;
  y = (*itr) - day - 1;

  score += c[after] * (z * (z + 1) / 2);
  score -= c[after] * ((x * (x + 1)) / 2 + (y * (y + 1)) / 2);
  return score;
}

void change(ll day, ll after) {
  ll prev = now[day];
  st[prev].erase(day);
  st[after].insert(day);
  now[day] = after;
  now_score = ch_score(day, after);
}
int main(int argc, char const *argv[]) {
  ll d;
  scanf("%lld", &d);
  init(d);
  re(i, 26) scanf("%lld", &c[i]);
  re(i, D) re(j, 26) scanf("%lld", &s[i][j]);
  // 初期値
  re(i, D) now[i] = i % 26;
  // 初期化
  make_set(now);
  now_score = calc(now);
  // 関数:
  //  change(day, after) dayをafterに変更
  //  ch_score(day, after) dayをafterに変更した場合の試算
  //  共に計算量10~20くらい
  //  シミュレート回数10^7
  ll sim = 20000000;
  while (--sim) {
    unsigned long long a = engine();
    unsigned long long b = engine();
    a %= D, b %= 26;
    if (now_score < ch_score(a, b)) {
      change(a, b);
    }
  }
  for (int i = 0; i < D; i++)
    std::cout << now[i] + 1 << '\\n';
  return 0;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <bitset>
#include <cmath>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <string>
#include <vector>
#define vll vector<ll>
#define vvvl vector<vvl>
#define vvl vector<vector<ll>>
#define VV(a, b, c, d) vector<vector<d>>(a, vector<d>(b, c))
#define VVV(a, b, c, d) vector<vvl>(a, vvl(b, vll(c, d)));
#define re(c, b) for (ll c = 0; c < b; c++)
#define all(obj) (obj).begin(), (obj).end()
typedef long long int ll;
typedef long double ld;
using namespace std;

#include <random>
std::random_device seed_gen;
std::mt19937_64 engine(seed_gen());
// ull a = engine();
ll D;
vll now, c(26);
vvl s;
vector<set<ll>> st(26);
ll now_score = 0;

void init(ll d) {
  D = d;
  now.resize(d, 0);
  s = VV(d, 26, 0, ll);
}
void make_set(vll t) {
  for (int i = 0; i < D; i++)
    st[t[i]].insert(i);
  for (int i = 0; i < 26; i++)
    st[i].insert(-1);
  for (int i = 0; i < 26; i++)
    st[i].insert(D);
}
ll calc(vll t) {
  ll score = 0;
  vll cnt(26, 0);
  for (int i = 0; i < D; i++) {
    ll idx = t[i];
    for (int j = 0; j < 26; j++) {
      if (j == idx)
        cnt[j] = 0;
      else {
        cnt[j]++;
        score -= cnt[j] * c[j];
      }
    }
    score += s[i][idx];
  }
  return score;
}

ll ch_score(ll day, ll after) {
  ll score = now_score;
  ll prev = now[day]; // 前に選択していたもの
  if (prev == after)
    return score;

  // prevを消したことで減るスコア
  score -= s[day][prev];

  auto itr = st[prev].find(day);
  auto p = itr, q = itr;
  ll x = day - (*--p) - 1;
  ll y = (*++q) - day - 1;

  ll z = x + y + 1;
  // x, y分増え、z分減る
  score += c[prev] * ((x * (x + 1)) / 2 + (y * (y + 1)) / 2);
  score -= c[prev] * ((z * (z + 1)) / 2);
  // afterを挿入した事で増えるスコア
  score += s[day][after];
  itr = st[after].upper_bound(day);
  p = itr;
  p--;

  z = (*itr) - (*p) - 1;
  x = day - (*p) - 1;
  y = (*itr) - day - 1;

  score += c[after] * (z * (z + 1) / 2);
  score -= c[after] * ((x * (x + 1)) / 2 + (y * (y + 1)) / 2);
  return score;
}

void change(ll day, ll after) {
  ll prev = now[day];
  st[prev].erase(day);
  st[after].insert(day);
  now[day] = after;
  now_score = ch_score(day, after);
}

int main(int argc, char const *argv[]) {
  ll d;
  std::cin >> d;
  init(d);
  re(i, 26) std::cin >> c[i];
  re(i, D) re(j, 26) std::cin >> s[i][j];

  // 初期値
  re(i, D) now[i] = i % 26;

  // 初期化
  make_set(now);
  now_score = calc(now);

  // 関数:
  //  change(day, after) dayをafterに変更
  //  ch_score(day, after) dayをafterに変更した場合の試算
  //  共に計算量10~20くらい
  //  シミュレート回数10^7

  ll sim = 23000000;

  while (--sim) {
    unsigned long long a = engine();
    unsigned long long b = engine();
    a %= D, b %= 26;
    if (now_score < ch_score(a, b)) {
      change(a, b);
    }
  }

  for (int i = 0; i < D; i++)
    std::cout << now[i] + 1 << '\\n';
  return 0;
}