Problem p02618 (C++) — buggy submission, verdict: Runtime Error

#include <bits/stdc++.h>
using namespace std;
const int N = 1005;
const long long inf = 6000000000000000000ll;
int Q, n, m, D, c[27], a[N][27], t[N], as[N], su[27], ans, i, j, x, y,
    s = -1000000000;
set<int> p[27];
int cal(int x, int y) {
  int sum = ans;
  sum -= a[x][t[x]];
  sum += a[x][y];
  auto it = p[t[x]].find(x);
  auto e = it;
  --e;
  ++it;
  sum +=
      c[t[x]] * (-(*it - *e) * (*it - *e - 1) / 2 +
                 (x - *e) * (x - *e - 1) / 2 + (*it - x) * (*it - x - 1) / 2);
  it = p[y].lower_bound(x);
  e = it;
  --e;
  sum -= c[y] * (-(*it - *e) * (*it - *e - 1) / 2 +
                 (x - *e) * (x - *e - 1) / 2 + (*it - x) * (*it - x - 1) / 2);
  return sum;
}
void Change(int x, int y) {
  ans -= a[x][t[x]];
  ans += a[x][y];
  auto it = p[t[x]].find(x);
  auto e = it;
  --e;
  ++it;
  ans +=
      c[t[x]] * (-(*it - *e) * (*it - *e - 1) / 2 +
                 (x - *e) * (x - *e - 1) / 2 + (*it - x) * (*it - x - 1) / 2);
  p[t[x]].erase(x);
  t[x] = y;
  p[t[x]].insert(x);
  it = p[t[x]].find(x);
  e = it;
  --e;
  ++it;
  ans -=
      c[t[x]] * (-(*it - *e) * (*it - *e - 1) / 2 +
                 (x - *e) * (x - *e - 1) / 2 + (*it - x) * (*it - x - 1) / 2);
}
void up() {
  if (s < ans) {
    s = ans;
    for (int i = 1; i <= D; ++i)
      as[i] = t[i];
  }
}
mt19937 rnd(time(0));
double rand_f() { return double(rnd()) / double(RAND_MAX); }
int u = 0;
void SA() {
  double T = 100000000;
  while (T > 0.001) {
    u = u % D + 1;
    int a = u, b = rnd() % 26 + 1;
    while (b == t[a])
      b = rnd() % 26 + 1;
    int tmp = cal(a, b);
    if (tmp > ans || rand_f() < exp((tmp - ans) / T)) {
      Change(a, b);
      up();
    }
    T = T * 0.9995;
  }
  // cout<<s<<' '<<ans<<endl;
}
void Find() {
  int i, j;
  for (i = 1; i <= 26; ++i) {
    p[i].clear();
    p[i].insert(D + 1);
    p[i].insert(0);
  }
  memset(su, 0, sizeof(su));
  ans = 0;
  for (i = 1; i <= D; ++i) {
    for (j = 1; j <= 26; ++j)
      ++su[j];
    t[i] = 1;
    for (j = 2; j <= 26; ++j)
      if (su[t[i]] * c[t[i]] + a[i][t[i]] < su[j] * c[j] + a[i][j])
        t[i] = j;
    ans += a[i][t[i]];
    su[t[i]] = 0;
    for (j = 1; j <= 26; ++j)
      ans -= su[j] * c[j];
    p[t[i]].insert(i);
    // as[i]=t[i];
  }
}
int main() {
  scanf("%d", &D);
  for (i = 1; i <= 26; ++i)
    scanf("%d", &c[i]);
  for (i = 1; i <= D; ++i)
    for (j = 1; j <= 26; ++j)
      scanf("%d", &a[i][j]);
  Find();
  up();
  while (1.0 * clock() / CLOCKS_PER_SEC < 1.8) {
    SA();
    SA();
    SA();
    SA();
    SA();
    SA();
    SA();
    SA();
    Find();
    up();
    // cout<<'?'<<' '<<s<<' '<<ans<<endl;
  }
  for (i = 1; i <= D; ++i) {
    cout << as[i] << endl;
    assert(as[i] >= 1 && as[i] <= 26);
  }
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
const int N = 1005;
const long long inf = 6000000000000000000ll;
int Q, n, m, D, c[27], a[N][27], t[N], su[27], ans, i, j, x, y, as[N], s;
set<int> p[27];
int cal(int x, int y) {
  int sum = ans;
  sum -= a[x][t[x]];
  sum += a[x][y];
  auto it = p[t[x]].find(x);
  auto e = it;
  --e;
  ++it;
  sum +=
      c[t[x]] * (-(*it - *e) * (*it - *e - 1) / 2 +
                 (x - *e) * (x - *e - 1) / 2 + (*it - x) * (*it - x - 1) / 2);
  it = p[y].lower_bound(x);
  e = it;
  --e;
  sum -= c[y] * (-(*it - *e) * (*it - *e - 1) / 2 +
                 (x - *e) * (x - *e - 1) / 2 + (*it - x) * (*it - x - 1) / 2);
  return sum;
}
void Change(int x, int y) {
  ans -= a[x][t[x]];
  ans += a[x][y];
  auto it = p[t[x]].find(x);
  auto e = it;
  --e;
  ++it;
  ans +=
      c[t[x]] * (-(*it - *e) * (*it - *e - 1) / 2 +
                 (x - *e) * (x - *e - 1) / 2 + (*it - x) * (*it - x - 1) / 2);
  p[t[x]].erase(x);
  t[x] = y;
  p[t[x]].insert(x);
  it = p[t[x]].find(x);
  e = it;
  --e;
  ++it;
  ans -=
      c[t[x]] * (-(*it - *e) * (*it - *e - 1) / 2 +
                 (x - *e) * (x - *e - 1) / 2 + (*it - x) * (*it - x - 1) / 2);
}
void up() {
  if (s < ans) {
    s = ans;
    for (int i = 1; i <= D; ++i)
      as[i] = t[i];
  }
}
mt19937 rnd(time(0));
double rand_f() { return double(rnd()) / double(RAND_MAX); }
int u = 0;
void SA() {
  double T = 100000000;
  while (T > 0.001) {
    u = u % D + 1;
    int a = u, b = rnd() % 26 + 1;
    while (b == t[a])
      b = rnd() % 26 + 1;
    int tmp = cal(a, b);
    if (tmp > ans || rand_f() < exp((tmp - ans) / T)) {
      Change(a, b);
      up();
    }
    T = T * 0.9995;
  }
  // cout<<s<<' '<<ans<<endl;
}
void Find() {
  for (i = 1; i <= 26; ++i) {
    p[i].clear();
    p[i].insert(D + 1);
    p[i].insert(0);
  }
  memset(su, 0, sizeof(su));
  ans = 0;
  int i, j;
  for (i = 1; i <= D; ++i) {
    for (j = 1; j <= 26; ++j)
      ++su[j];
    t[i] = 1;
    for (j = 2; j <= 26; ++j)
      if (su[t[i]] * c[t[i]] + a[i][t[i]] < su[j] * c[j] + a[i][j])
        t[i] = j;
    ans += a[i][t[i]];
    su[t[i]] = 0;
    for (j = 1; j <= 26; ++j)
      ans -= su[j] * c[j];
    p[t[i]].insert(i);
    // as[i]=t[i];
  }
}
int main() {
  scanf("%d", &D);
  for (i = 1; i <= 26; ++i) {
    scanf("%d", &c[i]);
    p[i].insert(D + 1);
    p[i].insert(0);
  }
  for (i = 1; i <= D; ++i)
    for (j = 1; j <= 26; ++j)
      scanf("%d", &a[i][j]);
  Find();
  up();
  while (1.0 * clock() / CLOCKS_PER_SEC < 1.8) {
    SA();
    SA();
    SA();
    SA();
    SA();
    SA();
    SA();
    SA();
    Find();
    up();
    // cout<<'?'<<' '<<ans<<endl;
  }
  for (i = 1; i <= D; ++i) {
    cout << as[i] << endl;
    assert(as[i] >= 1 && as[i] <= 26);
  }
}