Problem p02618 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>

#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tag_and_trait.hpp>
#include <ext/pb_ds/tree_policy.hpp>

using namespace std;
using ll = long long;

vector<int> c;
vector<vector<int>> s;
int D;
vector<int> v;
vector<set<int>> U;
vector<set<int>> minusU;
int now;

ll calc() {
  ll ans = 0;
  vector<int> last(26, 0);
  for (int i = 0; i < D; i++) {
    last[v[i] - 1] = i + 1;
    ans += s[i][v[i] - 1];
    // cout<<s[i][v[i] - 1]<<endl;
    for (int j = 0; j < 26; j++) {
      ans -= (ll)(i + 1 - last[j]) * c[j];
    }
  }
  now = ans;
  return ans;
}

ll calc(int d, int i) {
  int prev = v[d];
  now += s[d][i];
  now -= s[d][prev];
  int n = *U[i].upper_bound(d);
  int p = -*minusU[i].upper_bound(-d);
  int term = n - p - 1;
  now += term * (term + 1) / 2 * c[i];
  now -= (n - d) * (n - d - 1) / 2 + (d - p - 1) * (d - p) / 2 * c[i];
  n = *U[prev].upper_bound(d);
  p = -*minusU[prev].upper_bound(-d);
  term = n - p - 1;
  now -= term * (term + 1) / 2 * c[prev];
  now += (n - d) * (n - d - 1) / 2 + (d - p - 1) * (d - p) / 2 * c[prev];
  v[d] = i;
  U[prev].erase(d);
  minusU[prev].erase(-d);
  U[d].insert(d);
  minusU[prev].insert(-d);
  return now;
}

int retrnd(int n) { // 乱数生成
  std::random_device rnd;
  std::mt19937 mt(rnd());
  std::uniform_int_distribution<> rand100(0, n);
  return rand100(mt);
}

int main() {
  cin >> D;
  for (int i = 0; i < 26; i++) {
    int a;
    cin >> a;
    c.push_back(a);
  }
  for (int i = 0; i < D; i++) {
    vector<int> a;
    for (int j = 0; j < 26; j++) {
      int b;
      cin >> b;
      a.push_back(b);
    }
    s.push_back(a);
  }
  U = vector<set<int>>(26);
  minusU = vector<set<int>>(26);
  for (int i = 0; i < 26; i++) {
    U[i].insert(0);
    U[i].insert(D);
    minusU[i].insert(0);
    minusU[i].insert(D);
  }
  for (int i = 0; i < D; i++) {
    v.push_back(retrnd(25) + 1);
    U[v.back() - 1].insert(i);
    U[v.back() - 1].insert(-i);
  }
  now = calc();
  int try_cnt = 5000;
  for (int cnt = 0; cnt < try_cnt; cnt++) {
    int rnd = retrnd(D - 1);
    for (int i = 1; i < 27; i++) {
      int p = now;
      int prev = v[rnd];
      v[rnd] = i;
      calc();
      if (now < p) {
        v[rnd] = prev;
        calc();
      }
    }
  }
  for (int i = 0; i < D; i++) {
    cout << v[i] << endl;
  }
}

Fixed submission:
#include <bits/stdc++.h>

#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tag_and_trait.hpp>
#include <ext/pb_ds/tree_policy.hpp>

using namespace std;
using ll = long long;

vector<int> c;
vector<vector<int>> s;
int D;
vector<int> v;
vector<set<int>> U;
vector<set<int>> minusU;
int now;

ll calc() {
  ll ans = 0;
  vector<int> last(26, 0);
  for (int i = 0; i < D; i++) {
    last[v[i] - 1] = i + 1;
    ans += s[i][v[i] - 1];
    // cout<<s[i][v[i] - 1]<<endl;
    for (int j = 0; j < 26; j++) {
      ans -= (ll)(i + 1 - last[j]) * c[j];
    }
  }
  now = ans;
  return ans;
}

ll calc(int d, int i) {
  int prev = v[d];
  now += s[d][i];
  now -= s[d][prev];
  int n = *U[i].upper_bound(d);
  int p = -*minusU[i].upper_bound(-d);
  int term = n - p - 1;
  now += term * (term + 1) / 2 * c[i];
  now -= (n - d) * (n - d - 1) / 2 + (d - p - 1) * (d - p) / 2 * c[i];
  n = *U[prev].upper_bound(d);
  p = -*minusU[prev].upper_bound(-d);
  term = n - p - 1;
  now -= term * (term + 1) / 2 * c[prev];
  now += (n - d) * (n - d - 1) / 2 + (d - p - 1) * (d - p) / 2 * c[prev];
  v[d] = i;
  U[prev].erase(d);
  minusU[prev].erase(-d);
  U[d].insert(d);
  minusU[prev].insert(-d);
  return now;
}

int retrnd(int n) { // 乱数生成
  std::random_device rnd;
  std::mt19937 mt(rnd());
  std::uniform_int_distribution<> rand100(0, n);
  return rand100(mt);
}

int main() {
  cin >> D;
  for (int i = 0; i < 26; i++) {
    int a;
    cin >> a;
    c.push_back(a);
  }
  for (int i = 0; i < D; i++) {
    vector<int> a;
    for (int j = 0; j < 26; j++) {
      int b;
      cin >> b;
      a.push_back(b);
    }
    s.push_back(a);
  }
  U = vector<set<int>>(26);
  minusU = vector<set<int>>(26);
  for (int i = 0; i < 26; i++) {
    U[i].insert(0);
    U[i].insert(D);
    minusU[i].insert(0);
    minusU[i].insert(D);
  }
  for (int i = 0; i < D; i++) {
    v.push_back(retrnd(25) + 1);
    U[v.back() - 1].insert(i);
    U[v.back() - 1].insert(-i);
  }
  now = calc();
  int try_cnt = 1000;
  for (int cnt = 0; cnt < try_cnt; cnt++) {
    int rnd = retrnd(D - 1);
    for (int i = 1; i < 27; i++) {
      int p = now;
      int prev = v[rnd];
      v[rnd] = i;
      calc();
      if (now < p) {
        v[rnd] = prev;
        calc();
      }
    }
  }
  for (int i = 0; i < D; i++) {
    cout << v[i] << endl;
  }
}