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

#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <functional>
#include <iterator>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <time.h>
#include <vector>
#pragma warning(disable : 4996)

typedef long long ll;
typedef unsigned long long ull;
#define MIN(a, b) ((a) > (b) ? (b) : (a))
#define MAX(a, b) ((a) < (b) ? (b) : (a))
#define LINF 9223300000000000000
#define LINF2 1223300000000000000
#define LINF3 1000000000000
#define INF 2140000000
const long long MOD = 1000000007;
// const long long MOD = 998244353;

using namespace std;

int D;
int c[26];
vector<int> s;

void input() {
  scanf("%d", &D);
  s.resize(D * 26);
  int i, j;
  for (i = 0; i < 26; i++) {
    scanf("%d", &c[i]);
  }
#define P(i, j) ((i)*26 + (j))
  for (i = 0; i < D; i++) {
    for (j = 0; j < 26; j++) {
      scanf("%d", &s[P(i, j)]);
    }
  }
  return;
}

void calc(const vector<int> &t, vector<int> &ans) {
  vector<int> last(26, -1);
  int sum = 0;
  int i, k;
  for (i = 0; i < D; i++) {
    sum += s[P(i, t[i])];
    last[t[i]] = i;
    for (k = 0; k < 26; k++) {
      sum -= c[k] * (i - last[k]);
    }
    ans[i] = sum;
  }
  return;
}

void sel0(vector<int> &t) { return; }

void sel(vector<int> &t) {
  vector<int> last(26, -1);
  int i, j, k;
  for (i = 0; i < D; i++) {
    int max = -INF;
    int maxj = -1;
    for (j = 0; j < 26; j++) {
      int sum0 = s[P(i, j)];
      for (k = 0; k < 26; k++) {
        if (k != j)
          sum0 -= c[k] * (i - last[k]);
      }
      if (max < sum0) {
        max = sum0;
        maxj = j;
      }
    }
    t[i] = maxj;
    last[maxj] = i;
  }
  return;
}

void query(vector<int> &t, vector<set<int>> &z, int d, int q, int &ans) {
  int diff0 = 0, diff1 = 0;
  if (t[d] != q) {
    {
      int id = t[d];
      auto it = z[id].lower_bound(d);
      auto it2 = it;
      int prev = -1, next = D;
      if (it2 != z[id].begin()) {
        it2--;
        prev = *it2;
      }
      auto it3 = it;
      it3++;
      if (it3 != z[id].end()) {
        next = *it3;
      }
      int k0 = d - prev, k1 = next - d, k2 = next - prev;
      diff0 = k2 * (k2 - 1) / 2 - k0 * (k0 - 1) / 2 - k1 * (k1 - 1) / 2;
    }

    {
      int id = q;
      auto it = z[id].lower_bound(d);
      auto it2 = it;
      int prev = -1, next = D;
      if (it2 != z[id].begin()) {
        it2--;
        prev = *it2;
      }
      if (it != z[id].end()) {
        next = *it;
      }
      int k0 = d - prev, k1 = next - d, k2 = next - prev;
      diff1 = k2 * (k2 - 1) / 2 - k0 * (k0 - 1) / 2 - k1 * (k1 - 1) / 2;
    }
  }
  int diff = -c[t[d]] * diff0 + c[q] * diff1 - s[P(d, t[d])] + s[P(d, q)];
  if (diff > 0) {
    ans += diff;
    z[t[d]].erase(d);
    z[q].insert(d);
    t[d] = q;
  }

  return;
}

int Dice(int num) { return rand() % num; }

void solve() {
  input();

  int i;
#if 0
    vector<int> t(D);
    for (i = 0; i < D; i++) {
        scanf("%d", &t[i]); t[i]--;
    }
#else
  vector<int> t(D);
  // sel(t);
  sel0(t);
#endif

  vector<int> ans0(D);
  calc(t, ans0);
  int ans = ans0[D - 1];

  vector<set<int>> z(26);
  for (i = 0; i < D; i++) {
    z[t[i]].insert(i);
  }

  srand((unsigned int)time(NULL));

  int m = 10000;
  for (i = 0; i < m; i++) {
    vector<int> ts = t;
    vector<set<int>> zs = z;
    int anss = ans;
    int j;
    for (j = 0; j < 4; j++) {
      int d = Dice(D), q = Dice(26);
      query(t, z, d, q, ans);
      // printf("%d\\n", ans);
    }
    if (anss <= ans) {
      t.swap(ts);
      z.swap(zs);
      ans = anss;
    }
  }

  m = 5000000;
  for (i = 0; i < m; i++) {
    int d = Dice(D), q = Dice(26);
    query(t, z, d, q, ans);
    // printf("%d\\n", ans);
  }

#if 1
  for (i = 0; i < D; i++) {
    printf("%d\\n", t[i] + 1);
  }
#endif

#ifdef _DEBUG
  {
    vector<int> ans0(D);
    calc(t, ans0);
    printf("%d\\n", ans0[D - 1]);
  }
#endif

  return;
}

int main(int argc, char *argv[]) {
#if 1
  solve();
#else
  int T;
  scanf("%d", &T);
  int t;
  for (t = 0; t < T; t++) {
    // printf("Case #%d: ", t+1);
    solve();
  }
#endif
  return 0;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <functional>
#include <iterator>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <time.h>
#include <vector>
#pragma warning(disable : 4996)

typedef long long ll;
typedef unsigned long long ull;
#define MIN(a, b) ((a) > (b) ? (b) : (a))
#define MAX(a, b) ((a) < (b) ? (b) : (a))
#define LINF 9223300000000000000
#define LINF2 1223300000000000000
#define LINF3 1000000000000
#define INF 2140000000
const long long MOD = 1000000007;
// const long long MOD = 998244353;

using namespace std;

int D;
int c[26];
vector<int> s;

void input() {
  scanf("%d", &D);
  s.resize(D * 26);
  int i, j;
  for (i = 0; i < 26; i++) {
    scanf("%d", &c[i]);
  }
#define P(i, j) ((i)*26 + (j))
  for (i = 0; i < D; i++) {
    for (j = 0; j < 26; j++) {
      scanf("%d", &s[P(i, j)]);
    }
  }
  return;
}

void calc(const vector<int> &t, vector<int> &ans) {
  vector<int> last(26, -1);
  int sum = 0;
  int i, k;
  for (i = 0; i < D; i++) {
    sum += s[P(i, t[i])];
    last[t[i]] = i;
    for (k = 0; k < 26; k++) {
      sum -= c[k] * (i - last[k]);
    }
    ans[i] = sum;
  }
  return;
}

void sel0(vector<int> &t) { return; }

void sel(vector<int> &t) {
  vector<int> last(26, -1);
  int i, j, k;
  for (i = 0; i < D; i++) {
    int max = -INF;
    int maxj = -1;
    for (j = 0; j < 26; j++) {
      int sum0 = s[P(i, j)];
      for (k = 0; k < 26; k++) {
        if (k != j)
          sum0 -= c[k] * (i - last[k]);
      }
      if (max < sum0) {
        max = sum0;
        maxj = j;
      }
    }
    t[i] = maxj;
    last[maxj] = i;
  }
  return;
}

void query(vector<int> &t, vector<set<int>> &z, int d, int q, int &ans) {
  int diff0 = 0, diff1 = 0;
  if (t[d] != q) {
    {
      int id = t[d];
      auto it = z[id].lower_bound(d);
      auto it2 = it;
      int prev = -1, next = D;
      if (it2 != z[id].begin()) {
        it2--;
        prev = *it2;
      }
      auto it3 = it;
      it3++;
      if (it3 != z[id].end()) {
        next = *it3;
      }
      int k0 = d - prev, k1 = next - d, k2 = next - prev;
      diff0 = k2 * (k2 - 1) / 2 - k0 * (k0 - 1) / 2 - k1 * (k1 - 1) / 2;
    }

    {
      int id = q;
      auto it = z[id].lower_bound(d);
      auto it2 = it;
      int prev = -1, next = D;
      if (it2 != z[id].begin()) {
        it2--;
        prev = *it2;
      }
      if (it != z[id].end()) {
        next = *it;
      }
      int k0 = d - prev, k1 = next - d, k2 = next - prev;
      diff1 = k2 * (k2 - 1) / 2 - k0 * (k0 - 1) / 2 - k1 * (k1 - 1) / 2;
    }
  }
  int diff = -c[t[d]] * diff0 + c[q] * diff1 - s[P(d, t[d])] + s[P(d, q)];
  if (diff > 0) {
    ans += diff;
    z[t[d]].erase(d);
    z[q].insert(d);
    t[d] = q;
  }

  return;
}

int Dice(int num) { return rand() % num; }

void solve() {
  input();

  int i;
#if 0
    vector<int> t(D);
    for (i = 0; i < D; i++) {
        scanf("%d", &t[i]); t[i]--;
    }
#else
  vector<int> t(D);
  // sel(t);
  sel0(t);
#endif

  vector<int> ans0(D);
  calc(t, ans0);
  int ans = ans0[D - 1];

  vector<set<int>> z(26);
  for (i = 0; i < D; i++) {
    z[t[i]].insert(i);
  }

  srand((unsigned int)time(NULL));

  int m = 500000;
  for (i = 0; i < m; i++) {
    vector<int> ts = t;
    vector<set<int>> zs = z;
    int anss = ans;
    int j;
    for (j = 0; j < 4; j++) {
      int d = Dice(D), q = Dice(26);
      query(t, z, d, q, ans);
      // printf("%d\\n", ans);
    }
    if (anss <= ans) {
      t.swap(ts);
      z.swap(zs);
      ans = anss;
    }
  }

  m = 5000000;
  for (i = 0; i < m; i++) {
    int d = Dice(D), q = Dice(26);
    query(t, z, d, q, ans);
    // printf("%d\\n", ans);
  }

#if 1
  for (i = 0; i < D; i++) {
    printf("%d\\n", t[i] + 1);
  }
#endif

#ifdef _DEBUG
  {
    vector<int> ans0(D);
    calc(t, ans0);
    printf("%d\\n", ans0[D - 1]);
  }
#endif

  return;
}

int main(int argc, char *argv[]) {
#if 1
  solve();
#else
  int T;
  scanf("%d", &T);
  int t;
  for (t = 0; t < T; t++) {
    // printf("Case #%d: ", t+1);
    solve();
  }
#endif
  return 0;
}