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

#include <bits/stdc++.h>
using namespace std;

// #define int long long
#define rep(i, n) for (int i = 0; i < (int)(n); i++)
#define reps(i, n) for (int i = 1; i <= (int)(n); i++)
#define all(x) (x).begin(), (x).end()
#define uniq(x) (x).erase(unique(all(x)), (x).end())
#define bit(n) (1LL << (n))
#define cdiv(a, b) (((a)-1) / (b) + 1)
#define dump(x) cerr << #x " = " << (x) << endl
using vint = vector<int>;
using vvint = vector<vint>;
using pint = pair<int, int>;
using vpint = vector<pint>;
template <typename T>
using priority_queue_rev = priority_queue<T, vector<T>, greater<T>>;
constexpr double PI = 3.1415926535897932384626433832795028;
constexpr int DY[8] = {0, 1, 0, -1, 1, 1, -1, -1};
constexpr int DX[8] = {1, 0, -1, 0, 1, -1, -1, 1};
int gcd(int a, int b) {
  while (b) {
    swap(a %= b, b);
  }
  return a;
}
int lcm(int a, int b) { return a / gcd(a, b) * b; }
template <typename T> void fin(T mes) {
  cout << mes << endl;
  exit(0);
}
template <typename T, typename U> bool chmax(T &a, const U &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <typename T, typename U> bool chmin(T &a, const U &b) {
  if (b < a) {
    a = b;
    return true;
  }
  return false;
}
template <typename T, typename U>
ostream &operator<<(ostream &os, const pair<T, U> &rhs) {
  os << "(" << rhs.first << ", " << rhs.second << ")";
  return os;
}
template <typename T> ostream &operator<<(ostream &os, const vector<T> &rhs) {
  os << "{";
  for (auto itr = rhs.begin(); itr != rhs.end(); itr++) {
    os << *itr << (next(itr) != rhs.end() ? ", " : "");
  }
  os << "}";
  return os;
}
struct setup {
  static constexpr int PREC = 20;
  setup() {
    cout << fixed << setprecision(PREC);
    cerr << fixed << setprecision(PREC);
  };
} setup;

long long xor64(long long range) {
  static unsigned long long x = 1145141919810ULL;
  x ^= x << 13;
  x ^= x >> 7;
  return (x ^= x << 17) % range;
}

struct timer {
  unsigned long long begin_cycle;
  double cycle_per_sec = 2.8e9;
  unsigned long long get_cycle() {
    unsigned low, high;
    __asm__ volatile("rdtsc" : "=a"(low), "=d"(high));
    return ((unsigned long long)low) | ((unsigned long long)high << 32);
  }
  double get_time() { return (get_cycle() - begin_cycle) / cycle_per_sec; }
  void init() { begin_cycle = get_cycle(); }
} timer;

void simulated_annealing() {
  timer.init();
  double T = 5.95;
  double start_temp = 0.01;
  double end_temp = 0.01;
  int it = 0, pre_score;
  double time;
  while ((time = timer.get_time()) < T) {
    it++;
    double temp = start_temp + (end_temp - start_temp) * time / T;
    int new_score;
    // スコアが高いほどいい
    double prob = exp((new_score - pre_score) / temp);
    if (prob < (double)xor64(10000000) / 10000000) {
      // rejected
    } else {
      // accepted
    }
  }
}

int D;
int c[33];
int s[440][33];
int A = 26;
struct state {
  int val = 0;
  vint last = vint(A + 1, 0);
  vint history{0};
  state() {}
};
state beam_search() {
  int T = D;
  int beam_width = 300;
  vector<state> pre_states = {state()};
  reps(i, T) {
    vector<state> next_states;
    rep(j, pre_states.size()) {
      state pre_state = pre_states[j];
      int penalty = 0;
      reps(k, A) { penalty += c[k] * (i - pre_state.last[k]); }
      reps(k, A) {
        state next_state = pre_state;
        next_state.val += s[i][k] + c[k] * (i - pre_state.last[k]) - penalty;
        next_state.last[k] = i;
        next_state.history.push_back(k);
        next_states.push_back(next_state);
      }
    }

    // pre_statesから次の状態を計算しnext_statesに突っ込む
    vector<int> idx(next_states.size());
    for (int i = 0; i < idx.size(); i++) {
      idx[i] = i;
    }
    nth_element(idx.begin(), idx.begin() + beam_width - 1, idx.end(),
                [&](int l, int r) {
                  // lのほうが優先度が高い時、true
                  return next_states[l].val > next_states[r].val;
                });
    pre_states.clear();
    for (int i = 0; i < min(beam_width, (int)idx.size()); i++) {
      pre_states.push_back(next_states[idx[i]]);
    }
  }
  state ret;
  ret.val = INT_MIN;
  for (int i = 0; i < pre_states.size(); i++) {
    if (ret.val < pre_states[i].val) {
      ret = pre_states[i];
    }
  }
  return ret;
}

int score(vint history) {
  vint last(A + 1);
  int ret = 0;
  reps(i, D) {
    ret += s[i][history[i]];
    last[history[i]] = i;
    reps(j, A) { ret -= c[j] * (i - last[j]); }
  }
  return ret;
}

int main() {
  cin >> D;
  reps(i, A) { cin >> c[i]; }
  reps(i, D) {
    reps(j, A) { cin >> s[i][j]; }
  }
  state res = beam_search();
  reps(i, D) { cout << res.history[i] << endl; }
  dump(res.val);
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;

// #define int long long
#define rep(i, n) for (int i = 0; i < (int)(n); i++)
#define reps(i, n) for (int i = 1; i <= (int)(n); i++)
#define all(x) (x).begin(), (x).end()
#define uniq(x) (x).erase(unique(all(x)), (x).end())
#define bit(n) (1LL << (n))
#define cdiv(a, b) (((a)-1) / (b) + 1)
#define dump(x) cerr << #x " = " << (x) << endl
using vint = vector<int>;
using vvint = vector<vint>;
using pint = pair<int, int>;
using vpint = vector<pint>;
template <typename T>
using priority_queue_rev = priority_queue<T, vector<T>, greater<T>>;
constexpr double PI = 3.1415926535897932384626433832795028;
constexpr int DY[8] = {0, 1, 0, -1, 1, 1, -1, -1};
constexpr int DX[8] = {1, 0, -1, 0, 1, -1, -1, 1};
int gcd(int a, int b) {
  while (b) {
    swap(a %= b, b);
  }
  return a;
}
int lcm(int a, int b) { return a / gcd(a, b) * b; }
template <typename T> void fin(T mes) {
  cout << mes << endl;
  exit(0);
}
template <typename T, typename U> bool chmax(T &a, const U &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <typename T, typename U> bool chmin(T &a, const U &b) {
  if (b < a) {
    a = b;
    return true;
  }
  return false;
}
template <typename T, typename U>
ostream &operator<<(ostream &os, const pair<T, U> &rhs) {
  os << "(" << rhs.first << ", " << rhs.second << ")";
  return os;
}
template <typename T> ostream &operator<<(ostream &os, const vector<T> &rhs) {
  os << "{";
  for (auto itr = rhs.begin(); itr != rhs.end(); itr++) {
    os << *itr << (next(itr) != rhs.end() ? ", " : "");
  }
  os << "}";
  return os;
}
struct setup {
  static constexpr int PREC = 20;
  setup() {
    cout << fixed << setprecision(PREC);
    cerr << fixed << setprecision(PREC);
  };
} setup;

long long xor64(long long range) {
  static unsigned long long x = 1145141919810ULL;
  x ^= x << 13;
  x ^= x >> 7;
  return (x ^= x << 17) % range;
}

struct timer {
  unsigned long long begin_cycle;
  double cycle_per_sec = 2.8e9;
  unsigned long long get_cycle() {
    unsigned low, high;
    __asm__ volatile("rdtsc" : "=a"(low), "=d"(high));
    return ((unsigned long long)low) | ((unsigned long long)high << 32);
  }
  double get_time() { return (get_cycle() - begin_cycle) / cycle_per_sec; }
  void init() { begin_cycle = get_cycle(); }
} timer;

void simulated_annealing() {
  timer.init();
  double T = 5.95;
  double start_temp = 0.01;
  double end_temp = 0.01;
  int it = 0, pre_score;
  double time;
  while ((time = timer.get_time()) < T) {
    it++;
    double temp = start_temp + (end_temp - start_temp) * time / T;
    int new_score;
    // スコアが高いほどいい
    double prob = exp((new_score - pre_score) / temp);
    if (prob < (double)xor64(10000000) / 10000000) {
      // rejected
    } else {
      // accepted
    }
  }
}

int D;
int c[33];
int s[440][33];
int A = 26;
struct state {
  int val = 0;
  vint last = vint(A + 1, 0);
  vint history{0};
  state() {}
};
state beam_search() {
  int T = D;
  int beam_width = 500;
  vector<state> pre_states = {state()};
  reps(i, T) {
    vector<state> next_states;
    rep(j, pre_states.size()) {
      state pre_state = pre_states[j];
      int penalty = 0;
      reps(k, A) { penalty += c[k] * (i - pre_state.last[k]); }
      reps(k, A) {
        state next_state = pre_state;
        next_state.val += s[i][k] + c[k] * (i - pre_state.last[k]) - penalty;
        next_state.last[k] = i;
        next_state.history.push_back(k);
        next_states.push_back(next_state);
      }
    }

    // pre_statesから次の状態を計算しnext_statesに突っ込む
    vector<int> idx(next_states.size());
    for (int i = 0; i < idx.size(); i++) {
      idx[i] = i;
    }
    nth_element(idx.begin(), idx.begin() + beam_width - 1, idx.end(),
                [&](int l, int r) {
                  // lのほうが優先度が高い時、true
                  return next_states[l].val > next_states[r].val;
                });
    pre_states.clear();
    for (int i = 0; i < min(beam_width, (int)idx.size()); i++) {
      pre_states.push_back(next_states[idx[i]]);
    }
  }
  state ret;
  ret.val = INT_MIN;
  for (int i = 0; i < pre_states.size(); i++) {
    if (ret.val < pre_states[i].val) {
      ret = pre_states[i];
    }
  }
  return ret;
}

int score(vint history) {
  vint last(A + 1);
  int ret = 0;
  reps(i, D) {
    ret += s[i][history[i]];
    last[history[i]] = i;
    reps(j, A) { ret -= c[j] * (i - last[j]); }
  }
  return ret;
}

int main() {
  cin >> D;
  reps(i, A) { cin >> c[i]; }
  reps(i, D) {
    reps(j, A) { cin >> s[i][j]; }
  }
  state res = beam_search();
  reps(i, D) { cout << res.history[i] << endl; }
  dump(res.val);
}