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

#pragma GCC target("avx")
#pragma GCC optimize("O3")
#pragma GCC optimize("unroll-loops")
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
const int D = 365;
const int beam = 230;
int c[26];
int s[D][26];
int ans[D];

struct State {
  ll val = 0;
  int out[D] = {0};
  int cnt[26] = {0};
};

bool operator<(const State &a, const State &b) { return a.val < b.val; }
bool operator>(const State &a, const State &b) { return a.val > b.val; }

void solve(int ca) {
  int waste;
  scanf("%d", &waste);
  for (int i = 0; i < 26; ++i) {
    scanf("%d", c + i);
  }
  for (auto &i : s) {
    for (int &j : i) {
      scanf("%d", &j);
    }
  }

  priority_queue<State> pq[2];
  pq[0].push(State());
  for (int i = 0; i < D; ++i) {
    int prev = i & 1, next = prev ^ 1;
    for (int j = 0; !pq[prev].empty() && j < beam; ++j) {
      auto state = pq[prev].top();
      pq[prev].pop();
      ll dec = 0;
      for (int k = 0; k < 26; ++k) {
        state.cnt[k]++;
        dec += (ll)c[k] * state.cnt[k];
      }
      for (int k = 0; k < 26; ++k) {
        const ll inc = (ll)c[k] * state.cnt[k];
        const int tmp = state.cnt[k];
        const int tmp_val = state.val;
        state.val = state.val - dec + inc + s[i][k];
        state.cnt[k] = 0;
        state.out[i] = k;
        pq[next].push(state);
        state.cnt[k] = tmp;
        state.val = tmp_val;
      }
    }
    pq[prev] = priority_queue<State>();
  }
  auto best = pq[D & 1].top().out;
  for (int i = 0; i < D; ++i) {
    ans[i] = best[i] + 1;
  }
  for (int i = 0; i < D; ++i) {
    printf("%d\\n", ans[i]);
  }
}

// #define MULTI
int main() {
#ifdef MULTI
  int N;
  cin >> N;
  for (int i = 0; i < N; ++i) {
    solve(i + 1);
  }
#else
  solve(0);
#endif
  return 0;
}


Error: TLE

Fixed submission:

#pragma GCC target("avx")
#pragma GCC optimize("O3")
#pragma GCC optimize("unroll-loops")
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
const int D = 365;
const int beam = 240;
int c[26];
int s[D][26];
int ans[D];

struct State {
  ll val = 0;
  int out[D] = {0};
  int cnt[26] = {0};
};

bool operator<(const State &a, const State &b) { return a.val < b.val; }
bool operator>(const State &a, const State &b) { return a.val > b.val; }

void solve(int ca) {
  int waste;
  scanf("%d", &waste);
  for (int i = 0; i < 26; ++i) {
    scanf("%d", c + i);
  }
  for (auto &i : s) {
    for (int &j : i) {
      scanf("%d", &j);
    }
  }

  priority_queue<State> pq[2];
  pq[0].push(State());
  for (int i = 0; i < D; ++i) {
    int prev = i & 1, next = prev ^ 1;
    for (int j = 0; !pq[prev].empty() && j < beam; ++j) {
      auto state = pq[prev].top();
      pq[prev].pop();
      ll dec = 0;
      for (int k = 0; k < 26; ++k) {
        state.cnt[k]++;
        dec += (ll)c[k] * state.cnt[k];
      }
      for (int k = 0; k < 26; ++k) {
        const ll inc = (ll)c[k] * state.cnt[k];
        const int tmp = state.cnt[k];
        const int tmp_val = state.val;
        state.val = state.val - dec + inc + s[i][k];
        state.cnt[k] = 0;
        state.out[i] = k;
        pq[next].push(state);
        state.cnt[k] = tmp;
        state.val = tmp_val;
      }
    }
    pq[prev] = priority_queue<State>();
  }
  auto best = pq[D & 1].top().out;
  for (int i = 0; i < D; ++i) {
    ans[i] = best[i] + 1;
  }
  for (int i = 0; i < D; ++i) {
    printf("%d\\n", ans[i]);
  }
}

// #define MULTI
int main() {
#ifdef MULTI
  int N;
  cin >> N;
  for (int i = 0; i < N; ++i) {
    solve(i + 1);
  }
#else
  solve(0);
#endif
  return 0;
}