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

#include <algorithm>
#include <iostream>
#include <queue>
#include <stack>
#include <utility>
#include <vector>

#define REP(i, n) for (int i = 0; i < n; i++)
#define FOR(i, m, n) for (int i = m; i < n; i++)
#define ALL(a) (a).begin(), (a).end()

#define INF 1000000000
#define BEAM 10

using namespace std;

class state {
public:
  int day;
  int previous_score;
  vector<int> previous;
  vector<int> types;

  state() {}

  state(const state &a)
      : day(a.day), previous_score(a.previous_score), previous(a.previous),
        types(a.types) {}
};

int _scoring(const vector<int> &c, const vector<vector<int>> &s, int day,
             int type, const vector<int> &previous) {
  int score = s[day][type];

  REP(i, 26) { score -= c[i] * (day + 1 - previous[i]); }

  return score;
}

int scoring(const vector<int> &c, const vector<vector<int>> &s,
            state &this_state) {
  this_state.previous[this_state.types.back()] = this_state.day + 1;
  return this_state.previous_score + _scoring(c, s, this_state.day,
                                              this_state.types.back(),
                                              this_state.previous);
}

using queue_type = pair<state, int>;
struct CustomCompare {
  bool operator()(const queue_type &a, const queue_type &b) {
    return a.second < b.second;
  }
};

int main() {
  int D;
  cin >> D;

  vector<int> c(26);
  REP(i, 26)
  cin >> c[i];

  vector<vector<int>> s(D, vector<int>(26, 0));
  REP(i, D)
  REP(j, 26)
  cin >> s[i][j];

  vector<state> initial;
  REP(i, 26) {
    state this_state;
    this_state.day = 0;
    this_state.types.push_back(i);
    this_state.previous_score = 0;
    this_state.previous = vector<int>(26, 0);
    initial.push_back(this_state);
  }

  priority_queue<queue_type, vector<queue_type>, CustomCompare> q1, q2;
  REP(i, D) {
    if (i == 0) {
      // use initial
      for (auto i : initial) {
        q2.push(make_pair(i, scoring(c, s, i)));
      }
    } else {
      // use queue
      for (int ii = 0; ii < BEAM && !q1.empty(); ++ii) {
        auto i = q1.top();

        REP(j, 26) {
          state new_state = i.first;
          new_state.day++;
          new_state.types.push_back(j);

          q2.push(make_pair(new_state, scoring(c, s, new_state)));
        }

        q1.pop();
      }
    }
    swap(q1, q2);

    q2 = decltype(q2)();
  }

  auto res = q1.top().first;

  for (auto t : res.types)
    cout << t + 1 << endl;

  return 0;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <iostream>
#include <queue>
#include <stack>
#include <utility>
#include <vector>

#define REP(i, n) for (int i = 0; i < n; i++)
#define FOR(i, m, n) for (int i = m; i < n; i++)
#define ALL(a) (a).begin(), (a).end()

#define INF 1000000000
#define BEAM 10

using namespace std;

class state {
public:
  int day;
  int type;
  int previous_score;
  vector<stack<int>> previous;
  vector<int> types;

  state() {}

  state(const state &a)
      : day(a.day), type(a.type), previous_score(a.previous_score),
        previous(a.previous), types(a.types) {}
};

int _scoring(const vector<int> &c, const vector<vector<int>> &s, int day,
             int type, const vector<stack<int>> &previous) {
  int score = s[day][type];

  REP(i, 26) { score -= c[i] * (day + 1 - previous[i].top()); }

  return score;
}

int scoring(const vector<int> &c, const vector<vector<int>> &s,
            state &this_state) {
  this_state.previous[this_state.type].push(this_state.day + 1);
  int score =
      this_state.previous_score +
      _scoring(c, s, this_state.day, this_state.type, this_state.previous);
  // this_state.previous[this_state.type].pop();

  return score;
}

using queue_type = pair<state, int>;
struct CustomCompare {
  bool operator()(const queue_type &a, const queue_type &b) {
    return a.second < b.second;
  }
};

int main() {
  int D;
  cin >> D;

  vector<int> c(26);
  REP(i, 26)
  cin >> c[i];

  vector<vector<int>> s(D, vector<int>(26, 0));
  REP(i, D)
  REP(j, 26)
  cin >> s[i][j];

  vector<state> initial;
  REP(i, 26) {
    vector<stack<int>> previous(26);
    REP(j, 26)
    previous[j].push(0);

    state this_state;
    this_state.day = 0;
    this_state.type = i;
    this_state.types.push_back(i);
    this_state.previous_score = 0;
    this_state.previous = previous;
    initial.push_back(this_state);
  }

  priority_queue<queue_type, vector<queue_type>, CustomCompare> q1, q2;
  REP(i, D) {
    if (i == 0) {
      // use initial
      for (auto i : initial) {
        q2.push(make_pair(i, scoring(c, s, i)));
      }
    } else {
      // use queue
      for (int ii = 0; ii < BEAM && !q1.empty(); ++ii) {
        auto i = q1.top();

        REP(j, 26) {
          state new_state = i.first;
          new_state.day++;
          new_state.type = j;
          new_state.types.push_back(j);

          q2.push(make_pair(new_state, scoring(c, s, new_state)));
        }

        q1.pop();
      }
    }
    swap(q1, q2);

    q2 = decltype(q2)();
  }

  auto res = q1.top().first;

  for (auto t : res.types)
    cout << t + 1 << endl;

  return 0;
}