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

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
const ll MOD = (ll)1e9 + 7;
// const ll MOD=(ll)998244353;
const ll INF = (ll)1e18 + 10;
const double PI = 2.0 * acos(0);
template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> bool chmin(T &a, const T &b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}
int dy[] = {-1, 1, 0, 0};
int dx[] = {0, 0, -1, 1};
const int D = 365;
const int beam = 100;
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 itr() {
  priority_queue<State> pq[2];
  pq[0].push(State());
  for (int i = 0; i < D; ++i) {
    int prev = i & 1, next = prev ^ 1;
    // cout<<prev<<" "<<next<<endl;
    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) { // 開催しなかったときのへる値
        // cout<< state.cnt[k]<<endl;
        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);
        // cout<<"called"<<endl;
        state.cnt[k] = tmp;
        state.val = tmp_val;
        // if(i==0)cout<<state.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;
  }
}

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);
    }
  }

  itr();
  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:

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
const ll MOD = (ll)1e9 + 7;
// const ll MOD=(ll)998244353;
const ll INF = (ll)1e18 + 10;
const double PI = 2.0 * acos(0);
template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> bool chmin(T &a, const T &b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}
int dy[] = {-1, 1, 0, 0};
int dx[] = {0, 0, -1, 1};
const int D = 365;
const int beam = 200;
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 itr() {
  priority_queue<State> pq[2];
  int init[D] = {0};
  pq[0].push(State());
  for (int i = 0; i < D; ++i) {
    int prev = i & 1, next = prev ^ 1;
    // cout<<prev<<" "<<next<<endl;
    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) { // 開催しなかったときのへる値
        // cout<< state.cnt[k]<<endl;
        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);
        // cout<<"called"<<endl;
        state.cnt[k] = tmp;
        state.val = tmp_val;
        // if(i==0)cout<<state.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;
  }
}

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);
    }
  }

  itr();
  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;
}