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

#include <bits/stdc++.h>

using namespace std;

using ll = long long;

constexpr char newl = '\\n';

constexpr double ticks_per_sec = 2800000000;
constexpr double ticks_per_sec_inv = 1.0 / ticks_per_sec;
inline double rdtsc() { // in seconds
  uint32_t lo, hi;
  asm volatile("rdtsc" : "=a"(lo), "=d"(hi));
  return (((uint64_t)hi << 32) | lo) * ticks_per_sec_inv;
}
constexpr double TLE = 2; // sec
constexpr double time_limit = TLE * 0.995;

unsigned int randxor() {
  static unsigned int x = 123456789, y = 362436069, z = 521288629, w = 88675123;
  unsigned int t;
  t = (x ^ (x << 11));
  x = y;
  y = z;
  z = w;
  return (w = (w ^ (w >> 19)) ^ (t ^ (t >> 8)));
}

template <typename hoge> void my_shuffle(vector<hoge> &v) {
  int sz = v.size();
  for (int i = sz - 1; i >= 1; i--) {
    swap(v[i], v[randxor() % (i + 1)]);
  }
}

// 定数
constexpr int CONTEST_NUM = 26;

// 入力
constexpr int D = 365;
array<int, CONTEST_NUM> c;
array<array<int, CONTEST_NUM>, D> s;

ll calcScore(const array<int, D> &t) {
  ll score = 0;
  array<int, CONTEST_NUM> lasts{};

  for (int i = 0; i < D; i++) {
    lasts[t[i]] = i + 1;
    score += s[i][t[i]];

    for (int j = 0; j < CONTEST_NUM; j++) {
      score -= c[j] * (i + 1 - lasts[j]);
    }
  }
  return score;
}

void update(int target_day, int new_contest, array<int, D> &t,
            array<set<int>, CONTEST_NUM> &days) {
  int old_contest = t[target_day - 1];
  days[old_contest].erase(target_day);
  days[new_contest].insert(target_day);
  t[target_day - 1] = new_contest;
}

// return: diff score
ll calcDiff(int target_day, int new_contest, array<int, D> &t,
            array<set<int>, CONTEST_NUM> &days) {
  int old_contest = t[target_day - 1];

  ll diff = s[target_day - 1][new_contest] - s[target_day - 1][old_contest];
  {
    auto it = days[new_contest].upper_bound(target_day);
    int nex = *it;
    int prv = *prev(it);

    int d1 = target_day - prv - 1;
    int d2 = nex - target_day - 1;
    int d3 = nex - prv - 1;

    diff -= c[new_contest] * d1 * (d1 + 1) / 2;
    diff -= c[new_contest] * d2 * (d2 + 1) / 2;
    diff += c[new_contest] * d3 * (d3 + 1) / 2;
  }

  {
    auto it = days[old_contest].find(target_day);
    int nex = *next(it);
    int prv = *prev(it);

    int d1 = target_day - prv - 1;
    int d2 = nex - target_day - 1;
    int d3 = nex - prv - 1;

    diff += c[old_contest] * d1 * (d1 + 1) / 2;
    diff += c[old_contest] * d2 * (d2 + 1) / 2;
    diff -= c[old_contest] * d3 * (d3 + 1) / 2;
  }
  return diff;
}

struct State {
  int score;
  array<int, D> t;
  array<int, CONTEST_NUM> lasts;

  State() {}

  bool operator<(const State &s) const { return score > s.score; }
};

ll init(array<int, D> &t, array<set<int>, CONTEST_NUM> &days) {
  ll score = 0;
  array<int, CONTEST_NUM> lasts{};

  for (int i = 0; i < D; i++) {
    ll max_diff = -1e15;
    ll pena_sum = 0;
    for (int j = 0; j < CONTEST_NUM; j++) {
      ll tmp = c[j] * (i + 1 - lasts[j]);
      pena_sum += tmp;
      if (s[i][j] + tmp > max_diff) {
        max_diff = s[i][j] + tmp;
        t[i] = j;
      }
    }
    max_diff -= pena_sum;
    score += max_diff;
    lasts[t[i]] = i + 1;
    days[t[i]].insert(i + 1);
  }
  return score;
}

ll init2(array<int, D> &t, array<set<int>, CONTEST_NUM> &days) {
  vector<State> pq;
  pq.emplace_back();

  constexpr size_t BEAM = 250;
  for (int i = 0; i < D; i++) {
    vector<State> npq;
    int sz = min(pq.size(), BEAM);
    if (pq.size() > BEAM)
      nth_element(pq.begin(), pq.begin() + BEAM - 1, pq.end());
    for (int ii = 0; ii < sz; ++ii) {
      State cur = pq[ii];

      ll pena_sum = 0;
      for (int j = 0; j < CONTEST_NUM; j++) {
        ll tmp = c[j] * (i + 1 - cur.lasts[j]);
        pena_sum += tmp;
      }

      for (int j = 0; j < CONTEST_NUM; j++) {
        ll tmp = c[j] * (i + 1 - cur.lasts[j]);
        cur.t[i] = j;
        int sc = cur.score;
        int memo = cur.lasts[j];
        cur.lasts[j] = i + 1;
        cur.score += s[i][j] + tmp - pena_sum;
        npq.push_back(cur);
        cur.lasts[j] = memo;
        cur.score = sc;
      }
    }

    pq = move(npq);
  }

  State best = *min_element(pq.begin(), pq.end());
  t = move(best.t);
  for (int i = 0; i < D; i++) {
    days[t[i]].insert(i + 1);
  }
  return best.score;
}

int main() {
  double clock_begin = rdtsc();

  cin.tie(nullptr);
  ios::sync_with_stdio(false);

  {
    int dummy;
    cin >> dummy;

    for (int i = 0; i < CONTEST_NUM; i++) {
      cin >> c[i];
    }

    for (int i = 0; i < D; i++) {
      for (int j = 0; j < CONTEST_NUM; j++) {
        cin >> s[i][j];
      }
    }
  }

  array<set<int>, CONTEST_NUM> days;
  for (int i = 0; i < CONTEST_NUM; i++) {
    days[i].insert(0);
    days[i].insert(D + 1);
  }

  array<int, D> t{};
  ll current_score = init2(t, days);
  // assert(current_score == calcScore(t));
  ll best_score = current_score;
  array<int, D> ans = t;

  /*
  {
      array<int, D> t2{};
      for (int i = 0; i < D; i++) {
          for (int j = 1; j < CONTEST_NUM; j++) {
              if (s[i][j] > s[i][t2[i]]) t2[i] = j;
          }
      }
      cerr << current_score << " " << calcScore(t2) << newl;
  }*/

  // 焼きなまし
  /*
  {
      constexpr double start_temp = 10000;
      constexpr double end_temp = 100;
      constexpr ll R = 10000;
      double current_time;
      int counter = 0;
      while ((current_time = rdtsc() - clock_begin) < time_limit) {
          ++counter;
          int target_day = randxor() % D + 1;
          int new_contest = randxor() % (CONTEST_NUM - 1);
          new_contest += (new_contest >= t[target_day - 1]);

          ll next_score = current_score + calcDiff(target_day, new_contest, t,
  days); double temp = start_temp + (end_temp - start_temp) * current_time /
  time_limit; double prob = exp((next_score - current_score) / temp); bool
  force_next = prob > (double)(randxor() % R) / R; if (next_score >
  current_score || force_next) { update(target_day, new_contest, t, days);
              current_score = next_score;

              if (current_score > best_score) {
                  ans = t;
                  best_score = current_score;
              }
          }
      }
  }*/

  for (int i = 0; i < D; i++) {
    cout << ans[i] + 1 << newl;
  }

  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;

using ll = long long;

constexpr char newl = '\\n';

constexpr double ticks_per_sec = 2800000000;
constexpr double ticks_per_sec_inv = 1.0 / ticks_per_sec;
inline double rdtsc() { // in seconds
  uint32_t lo, hi;
  asm volatile("rdtsc" : "=a"(lo), "=d"(hi));
  return (((uint64_t)hi << 32) | lo) * ticks_per_sec_inv;
}
constexpr double TLE = 2; // sec
constexpr double time_limit = TLE * 0.995;

unsigned int randxor() {
  static unsigned int x = 123456789, y = 362436069, z = 521288629, w = 88675123;
  unsigned int t;
  t = (x ^ (x << 11));
  x = y;
  y = z;
  z = w;
  return (w = (w ^ (w >> 19)) ^ (t ^ (t >> 8)));
}

template <typename hoge> void my_shuffle(vector<hoge> &v) {
  int sz = v.size();
  for (int i = sz - 1; i >= 1; i--) {
    swap(v[i], v[randxor() % (i + 1)]);
  }
}

// 定数
constexpr int CONTEST_NUM = 26;

// 入力
constexpr int D = 365;
array<int, CONTEST_NUM> c;
array<array<int, CONTEST_NUM>, D> s;

ll calcScore(const array<int, D> &t) {
  ll score = 0;
  array<int, CONTEST_NUM> lasts{};

  for (int i = 0; i < D; i++) {
    lasts[t[i]] = i + 1;
    score += s[i][t[i]];

    for (int j = 0; j < CONTEST_NUM; j++) {
      score -= c[j] * (i + 1 - lasts[j]);
    }
  }
  return score;
}

void update(int target_day, int new_contest, array<int, D> &t,
            array<set<int>, CONTEST_NUM> &days) {
  int old_contest = t[target_day - 1];
  days[old_contest].erase(target_day);
  days[new_contest].insert(target_day);
  t[target_day - 1] = new_contest;
}

// return: diff score
ll calcDiff(int target_day, int new_contest, array<int, D> &t,
            array<set<int>, CONTEST_NUM> &days) {
  int old_contest = t[target_day - 1];

  ll diff = s[target_day - 1][new_contest] - s[target_day - 1][old_contest];
  {
    auto it = days[new_contest].upper_bound(target_day);
    int nex = *it;
    int prv = *prev(it);

    int d1 = target_day - prv - 1;
    int d2 = nex - target_day - 1;
    int d3 = nex - prv - 1;

    diff -= c[new_contest] * d1 * (d1 + 1) / 2;
    diff -= c[new_contest] * d2 * (d2 + 1) / 2;
    diff += c[new_contest] * d3 * (d3 + 1) / 2;
  }

  {
    auto it = days[old_contest].find(target_day);
    int nex = *next(it);
    int prv = *prev(it);

    int d1 = target_day - prv - 1;
    int d2 = nex - target_day - 1;
    int d3 = nex - prv - 1;

    diff += c[old_contest] * d1 * (d1 + 1) / 2;
    diff += c[old_contest] * d2 * (d2 + 1) / 2;
    diff -= c[old_contest] * d3 * (d3 + 1) / 2;
  }
  return diff;
}

struct State {
  int score;
  array<int, D> t;
  array<int, CONTEST_NUM> lasts;

  State() {}

  bool operator<(const State &s) const { return score > s.score; }
};

ll init(array<int, D> &t, array<set<int>, CONTEST_NUM> &days) {
  ll score = 0;
  array<int, CONTEST_NUM> lasts{};

  for (int i = 0; i < D; i++) {
    ll max_diff = -1e15;
    ll pena_sum = 0;
    for (int j = 0; j < CONTEST_NUM; j++) {
      ll tmp = c[j] * (i + 1 - lasts[j]);
      pena_sum += tmp;
      if (s[i][j] + tmp > max_diff) {
        max_diff = s[i][j] + tmp;
        t[i] = j;
      }
    }
    max_diff -= pena_sum;
    score += max_diff;
    lasts[t[i]] = i + 1;
    days[t[i]].insert(i + 1);
  }
  return score;
}

ll init2(array<int, D> &t, array<set<int>, CONTEST_NUM> &days) {
  vector<State> pq;
  pq.emplace_back();

  constexpr size_t BEAM = 300;
  for (int i = 0; i < D; i++) {
    vector<State> npq;
    int sz = min(pq.size(), BEAM);
    if (pq.size() > BEAM)
      nth_element(pq.begin(), pq.begin() + BEAM - 1, pq.end());
    for (int ii = 0; ii < sz; ++ii) {
      State cur = pq[ii];

      ll pena_sum = 0;
      for (int j = 0; j < CONTEST_NUM; j++) {
        ll tmp = c[j] * (i + 1 - cur.lasts[j]);
        pena_sum += tmp;
      }

      for (int j = 0; j < CONTEST_NUM; j++) {
        ll tmp = c[j] * (i + 1 - cur.lasts[j]);
        cur.t[i] = j;
        int sc = cur.score;
        int memo = cur.lasts[j];
        cur.lasts[j] = i + 1;
        cur.score += s[i][j] + tmp - pena_sum;
        npq.push_back(cur);
        cur.lasts[j] = memo;
        cur.score = sc;
      }
    }

    pq = move(npq);
  }

  State best = *min_element(pq.begin(), pq.end());
  t = move(best.t);
  for (int i = 0; i < D; i++) {
    days[t[i]].insert(i + 1);
  }
  return best.score;
}

int main() {
  double clock_begin = rdtsc();

  cin.tie(nullptr);
  ios::sync_with_stdio(false);

  {
    int dummy;
    cin >> dummy;

    for (int i = 0; i < CONTEST_NUM; i++) {
      cin >> c[i];
    }

    for (int i = 0; i < D; i++) {
      for (int j = 0; j < CONTEST_NUM; j++) {
        cin >> s[i][j];
      }
    }
  }

  array<set<int>, CONTEST_NUM> days;
  for (int i = 0; i < CONTEST_NUM; i++) {
    days[i].insert(0);
    days[i].insert(D + 1);
  }

  array<int, D> t{};
  ll current_score = init2(t, days);
  // assert(current_score == calcScore(t));
  ll best_score = current_score;
  array<int, D> ans = t;

  /*
  {
      array<int, D> t2{};
      for (int i = 0; i < D; i++) {
          for (int j = 1; j < CONTEST_NUM; j++) {
              if (s[i][j] > s[i][t2[i]]) t2[i] = j;
          }
      }
      cerr << current_score << " " << calcScore(t2) << newl;
  }*/

  // 焼きなまし
  /*
  {
      constexpr double start_temp = 10000;
      constexpr double end_temp = 100;
      constexpr ll R = 10000;
      double current_time;
      int counter = 0;
      while ((current_time = rdtsc() - clock_begin) < time_limit) {
          ++counter;
          int target_day = randxor() % D + 1;
          int new_contest = randxor() % (CONTEST_NUM - 1);
          new_contest += (new_contest >= t[target_day - 1]);

          ll next_score = current_score + calcDiff(target_day, new_contest, t,
  days); double temp = start_temp + (end_temp - start_temp) * current_time /
  time_limit; double prob = exp((next_score - current_score) / temp); bool
  force_next = prob > (double)(randxor() % R) / R; if (next_score >
  current_score || force_next) { update(target_day, new_contest, t, days);
              current_score = next_score;

              if (current_score > best_score) {
                  ans = t;
                  best_score = current_score;
              }
          }
      }
  }*/

  for (int i = 0; i < D; i++) {
    cout << ans[i] + 1 << newl;
  }

  return 0;
}