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

#include <algorithm>
#include <array>
#include <bitset>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <functional>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <string>
#include <sys/time.h>
#include <tuple>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

using namespace std;
using ll = long long;
using ull = unsigned long long;
using ld = long double;
#define REP(i, m, n) for (int i = (m); i < (n); ++i)
#define rep(i, n) REP(i, 0, n)
#define all(s) (s).begin(), (s).end()
#define pb push_back
#define fs first
#define sc second
template <typename T> bool chmax(T &a, const T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <typename T> bool chmin(T &a, const T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
template <typename T> ll pow2(const T n) { return (1LL << n); }
void setp(const int n) { cout << fixed << setprecision(n); }
constexpr int INF = 1e9 + 1;
constexpr ll LINF = 1e18 + 1;
constexpr ll MOD = 1e9 + 7;
// constexpr ll MOD = 998244353;
constexpr ld EPS = 1e-11;
const ld PI = acos(-1);

constexpr int BEAM_SIZE = 1000;
constexpr int SZ = 26;
constexpr int D = 365;
array<int, SZ> C;
array<array<int, SZ>, D> S;

struct State {
  array<int, SZ> last = {};
  array<int, D> ans;
  int score = 0;

  void update_score(int d, int t) {
    score += S[d][t];
    rep(i, SZ) score -= C[i] * (d + 1 - last[i]);
  }

  int next_eval(int d, int t) {
    int res = score;
    res += S[d][t];
    rep(i, SZ) {
      if (i != t)
        res -= C[i] * (d + 1 - last[i]);
    }
    return res;
  }

  int eval() const { return score; }

  State place(int d, int t) {
    State res = *this;
    res.ans[d] = t;
    res.last[t] = d + 1;
    res.update_score(d, t);
    return res;
  }

  array<int, D> answer() { return ans; }
};

array<int, D> solve() {
  State initial_state;
  using f = function<bool(State &, State &)>;
  priority_queue<State, vector<State>, f> cur(
      [](State &a, State &b) { return a.eval() > b.eval(); });
  cur.push(initial_state);
  priority_queue<State, vector<State>, f> nxt(
      [](State &a, State &b) { return a.eval() > b.eval(); });
  rep(d, D) {
    while (!nxt.empty())
      nxt.pop();
    while (!cur.empty()) {
      State now = cur.top();
      cur.pop();
      rep(t, SZ) {
        if (nxt.size() < BEAM_SIZE || now.next_eval(d, t) > nxt.top().eval())
          nxt.push(now.place(d, t));
        if (nxt.size() > BEAM_SIZE)
          nxt.pop();
      }
    }
    swap(cur, nxt);
  }

  State res;
  while (!cur.empty()) {
    if (cur.size() == 1)
      res = cur.top();
    cur.pop();
  }
  cerr << res.eval() << endl;
  return res.answer();
}

int main(void) {
  int tmp;
  cin >> tmp;
  rep(i, SZ) cin >> C[i];
  rep(i, D) rep(j, SZ) cin >> S[i][j];

  auto ans = solve();
  for (int i : ans)
    cout << i + 1 << endl;

  return 0;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <array>
#include <bitset>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <functional>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <string>
#include <sys/time.h>
#include <tuple>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

using namespace std;
using ll = long long;
using ull = unsigned long long;
using ld = long double;
#define REP(i, m, n) for (int i = (m); i < (n); ++i)
#define rep(i, n) REP(i, 0, n)
#define all(s) (s).begin(), (s).end()
#define pb push_back
#define fs first
#define sc second
template <typename T> bool chmax(T &a, const T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <typename T> bool chmin(T &a, const T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
template <typename T> ll pow2(const T n) { return (1LL << n); }
void setp(const int n) { cout << fixed << setprecision(n); }
constexpr int INF = 1e9 + 1;
constexpr ll LINF = 1e18 + 1;
constexpr ll MOD = 1e9 + 7;
// constexpr ll MOD = 998244353;
constexpr ld EPS = 1e-11;
const ld PI = acos(-1);

constexpr int BEAM_SIZE = 1100;
constexpr int SZ = 26;
constexpr int D = 365;
array<int, SZ> C;
array<array<int, SZ>, D> S;

struct State {
  array<int, SZ> last = {};
  array<int, D> ans;
  int score = 0;

  void update_score(int d, int t) {
    score += S[d][t];
    rep(i, SZ) score -= C[i] * (d + 1 - last[i]);
  }

  int next_eval(int d, int t) {
    int res = score;
    res += S[d][t];
    rep(i, SZ) {
      if (i != t)
        res -= C[i] * (d + 1 - last[i]);
    }
    return res;
  }

  int eval() const { return score; }

  State place(int d, int t) {
    State res = *this;
    res.ans[d] = t;
    res.last[t] = d + 1;
    res.update_score(d, t);
    return res;
  }

  array<int, D> answer() { return ans; }
};

array<int, D> solve() {
  State initial_state;
  using f = function<bool(State &, State &)>;
  priority_queue<State, vector<State>, f> cur(
      [](State &a, State &b) { return a.eval() > b.eval(); });
  cur.push(initial_state);
  priority_queue<State, vector<State>, f> nxt(
      [](State &a, State &b) { return a.eval() > b.eval(); });
  rep(d, D) {
    while (!nxt.empty())
      nxt.pop();
    while (!cur.empty()) {
      State now = cur.top();
      cur.pop();
      rep(t, SZ) {
        if (nxt.size() < BEAM_SIZE || now.next_eval(d, t) > nxt.top().eval())
          nxt.push(now.place(d, t));
        if (nxt.size() > BEAM_SIZE)
          nxt.pop();
      }
    }
    swap(cur, nxt);
  }

  State res;
  while (!cur.empty()) {
    if (cur.size() == 1)
      res = cur.top();
    cur.pop();
  }
  cerr << res.eval() << endl;
  return res.answer();
}

int main(void) {
  int tmp;
  cin >> tmp;
  rep(i, SZ) cin >> C[i];
  rep(i, D) rep(j, SZ) cin >> S[i][j];

  auto ans = solve();
  for (int i : ans)
    cout << i + 1 << endl;

  return 0;
}