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


#include <bits/stdc++.h>
using namespace std;
#define rep(i, n) for (int i = 0; i < int(n); ++i)
#define all(a) (a).begin(), (a).end()
typedef long long ll;
#define rep1(i, n) for (int i = 1; i <= int(n); ++i)
#define lb(s, x) (lower_bound(all(s), x) - s.begin())
#define ub(s, x) (upper_bound(all(s), x) - s.begin())
#define uniq(a) a.erase(unique(all(a)), a.end())
#define bit(k) (1LL << (k))
const int INF = 1e9;

#define debug(...) recursive_debug(#__VA_ARGS__, __VA_ARGS__)

void recursive_debug(string) { cerr << endl; }

template <class H, class... T> void recursive_debug(string s, H h, T... t) {
  int comma = s.find(',');
  cerr << s.substr(0, comma) << " = " << h << ", ";
  recursive_debug(s.substr(comma + 1), t...);
}

template <class T, class U>
ostream &operator<<(ostream &ost, const pair<T, U> &p) {
  ost << "(" << p.first << ", " << p.second << ")";
  return ost;
}

template <class T> ostream &operator<<(ostream &ost, const vector<T> &v) {
  ost << "[";
  for (T t : v)
    ost << t << ",";
  ost << "]";
  return ost;
}

struct State {
  ll score;
  vector<int> last;
  vector<int> types;
  bool operator<(const State &rhs) const { return (score < rhs.score); }
  bool operator>(const State &rhs) const { return (score > rhs.score); }
};

int main(void) {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);
  int D, M = 26;
  cin >> D;
  vector<ll> c(M);
  rep(m, M) cin >> c[m];
  vector<vector<ll>> s(D, vector<ll>(26));
  rep(d, D) rep(m, M) cin >> s[d][m];

  set<State> cur_states = {{0, vector<int>(M, -1), {}}};
  size_t SIZE = 840;

  rep(d, D) {
    set<State> nxt_states;
    for (auto &cur_itr : cur_states) {
      State state = cur_itr;
      rep(m, M) {
        State nxt_state;
        nxt_state.score = state.score + s[d][m];
        nxt_state.last = state.last;
        nxt_state.last[m] = d;
        nxt_state.types = state.types;
        nxt_state.types.push_back(m);

        rep(i, M) { nxt_state.score -= c[i] * (d - nxt_state.last[i]); }
        nxt_states.insert(nxt_state);
        if (nxt_states.size() > SIZE) {
          nxt_states.erase(nxt_states.begin());
        }
      }
    }

    cur_states = std::move(nxt_states);
  }

  for (int t : cur_states.begin()->types) {
    cout << t + 1 << endl;
  }
  // debug(cur_states[0].score);

  return 0;
}


Error: TLE

Fixed submission:


#include <bits/stdc++.h>
using namespace std;
#define rep(i, n) for (int i = 0; i < int(n); ++i)
#define all(a) (a).begin(), (a).end()
typedef long long ll;
#define rep1(i, n) for (int i = 1; i <= int(n); ++i)
#define lb(s, x) (lower_bound(all(s), x) - s.begin())
#define ub(s, x) (upper_bound(all(s), x) - s.begin())
#define uniq(a) a.erase(unique(all(a)), a.end())
#define bit(k) (1LL << (k))
const int INF = 1e9;

#define debug(...) recursive_debug(#__VA_ARGS__, __VA_ARGS__)

void recursive_debug(string) { cerr << endl; }

template <class H, class... T> void recursive_debug(string s, H h, T... t) {
  int comma = s.find(',');
  cerr << s.substr(0, comma) << " = " << h << ", ";
  recursive_debug(s.substr(comma + 1), t...);
}

template <class T, class U>
ostream &operator<<(ostream &ost, const pair<T, U> &p) {
  ost << "(" << p.first << ", " << p.second << ")";
  return ost;
}

template <class T> ostream &operator<<(ostream &ost, const vector<T> &v) {
  ost << "[";
  for (T t : v)
    ost << t << ",";
  ost << "]";
  return ost;
}

struct State {
  ll score;
  vector<int> last;
  vector<int> types;
  bool operator<(const State &rhs) const { return (score < rhs.score); }
  bool operator>(const State &rhs) const { return (score > rhs.score); }
};

int main(void) {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);
  int D, M = 26;
  cin >> D;
  vector<ll> c(M);
  rep(m, M) cin >> c[m];
  vector<vector<ll>> s(D, vector<ll>(26));
  rep(d, D) rep(m, M) cin >> s[d][m];

  priority_queue<State, vector<State>, greater<State>> cur_states;
  cur_states.push({0, vector<int>(M, -1), {}});
  size_t SIZE = 845;

  rep(d, D) {
    priority_queue<State, vector<State>, greater<State>> nxt_states;
    while (!cur_states.empty()) {
      State state = cur_states.top();
      cur_states.pop();
      rep(m, M) {
        State nxt_state;
        nxt_state.score = state.score + s[d][m];
        nxt_state.last = state.last;
        nxt_state.last[m] = d;
        nxt_state.types = state.types;
        nxt_state.types.push_back(m);

        rep(i, M) { nxt_state.score -= c[i] * (d - nxt_state.last[i]); }
        nxt_states.push(nxt_state);
        if (nxt_states.size() > SIZE) {
          nxt_states.pop();
        }
      }
    }

    cur_states = std::move(nxt_states);
  }
  while ((int)cur_states.size() > 1) {
    cur_states.pop();
  }
  vector<int> types = cur_states.top().types;
  for (int t : types) {
    cout << t + 1 << endl;
  }
  // debug(cur_states.top().score);

  return 0;
}