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

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using ull = unsigned long long;
#define REP(i, n) for (int i = 0; i < (n); ++i)
#define RREP(i, n) for (int i = (n)-1; i >= 0; --i)
#define ALL(v) (v).begin(), (v).end()
template <class T>
using reverse_priority_queue = priority_queue<T, vector<T>, greater<T>>;

const int D = 365;
const int NTYPE = 26;
const int MAX_DELTA_DAY = 10;

minstd_rand rng;
clock_t start;

array<int, NTYPE> c;
array<array<int, NTYPE>, D> s;

const int NUM_CAND_TYPE = NTYPE / 2;
array<array<int, NUM_CAND_TYPE>, D> cand_types;
array<set<int>, D> cand_types_set;

array<int, D> day_order;
array<int, NTYPE> type_order;
array<int, NUM_CAND_TYPE> type_order2;

struct TypeSchedule {
  int id;
  set<int> day;
  int score;

  int get_prev_day(int x) const {
    auto it1 = day.lower_bound(x);
    int prev_day = -1;
    if (it1 != day.begin()) {
      prev_day = *(--it1);
    }
    return prev_day;
  }

  int get_next_day(int x) const {
    auto it2 = day.upper_bound(x);
    if (it2 != day.end() && *it2 == x) {
      ++it2;
    }
    int next_day = D;
    if (it2 != day.end()) {
      next_day = *it2;
    }
    return next_day;
  }

  TypeSchedule(int id) : id(id) { score = -c.at(id) * (D + 1) * D / 2; }

  int add_gain(int x) const {
    const int prev_day = get_prev_day(x);
    const int next_day = get_next_day(x);
    // prev_day + 1, ... , next_day - 1 = next_day - prev_day - 1 個
    const int l = next_day - prev_day - 1;
    int g = c.at(id) * (l + 1) * l / 2;
    // prev_day + 1, ..., x - 1 = x - prev_day - 1 個
    const int m1 = x - prev_day - 1;
    const int m2 = next_day - x - 1;
    g -= c.at(id) * (m1 * (m1 + 1) / 2 + m2 * (m2 + 1) / 2);
    g += s.at(x).at(id);
    return g;
  }

  void add(int x) {
    score += add_gain(x);
    day.insert(x);
  }

  int remove_gain(int x) const {
    assert(day.count(x) == 1);
    const int prev_day = get_prev_day(x);
    const int next_day = get_next_day(x);
    // prev_day + 1, ... , next_day - 1 = next_day - prev_day - 1 個
    const int l = next_day - prev_day - 1;
    int g = -c.at(id) * (l + 1) * l / 2;
    // prev_day + 1, ..., x - 1 = x - prev_day - 1 個
    const int m1 = x - prev_day - 1;
    const int m2 = next_day - x - 1;
    g += c.at(id) * (m1 * (m1 + 1) / 2 + m2 * (m2 + 1) / 2);
    g -= s.at(x).at(id);
    return g;
  }

  void remove(int x) {
    score += remove_gain(x);
    day.erase(x);
  }
};

struct Schedule {
  vector<TypeSchedule> type;
  array<int, D> d2t;
  int score;
  Schedule() {
    fill(ALL(d2t), -1);
    score = 0;
    REP(i, NTYPE) {
      type.emplace_back(i);
      score += type.at(i).score;
    }
    validate_score();
  }

  int calc_gain_set(int d, int t) const {
    const int t0 = d2t.at(d);
    if (t0 == t) {
      return 0;
    }
    int ret = 0;
    if (t0 != -1) {
      ret += type.at(t0).remove_gain(d);
    }
    ret += type.at(t).add_gain(d);
    return ret;
  }

  void set(int d, int t) {
    if (d2t.at(d) == t) {
      return;
    }
    score += calc_gain_set(d, t);
    const int t0 = d2t.at(d);
    d2t.at(d) = t;
    type.at(t).add(d);
    if (t0 != -1) {
      type.at(t0).remove(d);
    }
    validate_score();
  }

  int get(int d) const { return d2t.at(d); }
  void unset(int d, int t) {
    score -= type.at(t).score;
    d2t.at(d) = -1;
    type.at(t).remove(d);
    score += type.at(t).score;
    validate_score();
  }
  void swap(int d1, int d2) {
    if (get(d1) == get(d2)) {
      return;
    }
    const int t1 = get(d1);
    const int t2 = get(d2);
    set(d1, t2);
    set(d2, t1);
    validate_score();
  }

  void output() const {
    REP(i, D) {
      assert(d2t.at(i) != -1);
      cout << d2t.at(i) + 1 << endl;
    }
  }

  void validate_score() const {
    return;
    // int tmp = 0;
    // array<int, NTYPE> last;
    // fill(ALL(last), 0);
    // REP(i, D){
    //     int t = get(i);
    //     if(t != -1){
    //         tmp += s.at(i).at(t);
    //         last.at(t) = i + 1;
    //     }
    //     REP(s, NTYPE){
    //         tmp -= c.at(s) * (i + 1 - last.at(s));
    //     }
    // }
    // assert(tmp == score);
  }
};

void input() {
  int _d;
  cin >> _d;
  REP(i, NTYPE) { cin >> c.at(i); }
  REP(i, D) {
    REP(j, NTYPE) { cin >> s.at(i).at(j); }
  }
}

int calc_score_type(const array<int, D> &sch, int t) {
  int last = 0;
  int ret = 0;
  REP(i, D) {
    if (sch.at(i) == t) {
      ret += s.at(i).at(t);
      last = i + 1;
    }
    ret -= c.at(t) * (i + 1 - last);
  }
  return ret;
}

int calc_satisfaction(const array<int, D> &sch) {
  int ret = 0;
  REP(i, NTYPE) { ret += calc_score_type(sch, i); }
  return ret;
}

void greedy_init(Schedule &ans, int k = 1) {
  array<int, NTYPE> last;
  fill(ALL(last), 0);
  REP(i, D) {
    int max_type = -1;
    int max_val = numeric_limits<int>::min();
    REP(j, NTYPE) {
      int tmp = s.at(i).at(j);
      tmp += c.at(j) * (i + 1 - last.at(j)) * (min(D, i + k + 1) - i);
      if (max_val < tmp) {
        max_type = j;
        max_val = tmp;
      }
    }
    assert(max_type != -1);
    ans.set(i, max_type);
    last.at(max_type) = i + 1;
  }
}

void init() {
  REP(i, D) { day_order.at(i) = i; }
  REP(i, NTYPE) { type_order.at(i) = i; }
  REP(i, NUM_CAND_TYPE) { type_order2.at(i) = i; }

  REP(i, D) {
    array<pair<int, int>, NTYPE> tmp;
    REP(j, NTYPE) { tmp.at(j) = {s.at(i).at(j) + c.at(j), j}; }
    sort(ALL(tmp), greater<pair<int, int>>());
    REP(j, NUM_CAND_TYPE) {
      cand_types.at(i).at(j) = tmp.at(j).second;
      cand_types_set.at(i).insert(tmp.at(j).second);
    }
  }
}

double get_cur_sec() { return (double)(clock() - start) / CLOCKS_PER_SEC; }

int calc_gain_swap(Schedule &ans, int i, int j) {
  if (ans.get(i) == ans.get(j)) {
    return 0;
  }
  const int ti = ans.get(i);
  const int tj = ans.get(j);
  assert(ti != -1 && tj != -1);
  int g = -ans.score;
  ans.set(i, tj);
  ans.set(j, ti);
  g += ans.score;
  ans.set(i, ti);
  ans.set(j, tj);
  return g;
}

void greedy_update(Schedule &ans) {
  bool modified = true;
  while (modified) {
    modified = false;
    // change type
    REP(j, D) {
      REP(k, NTYPE) {
        if (ans.get(j) == k) {
          continue;
        }
        const int g = ans.calc_gain_set(j, k);
        if (g > 0) {
          modified = true;
          ans.set(j, k);
        }
      }
    }

    // swap day
    REP(j, D) {
      for (int delta = 1; delta <= MAX_DELTA_DAY && j + delta < D; ++delta) {
        const int g = calc_gain_swap(ans, j, j + delta);
        if (g > 0) {
          modified = true;
          ans.swap(j, j + delta);
        }
      }
    }

    if (get_cur_sec() >= 1.9) {
      break;
    }
  }
}

void search_cycle(Schedule &out, double exit_time, int max_it) {
  for (int it = 0; it < max_it; ++it) {
    REP(j, D) {
      Schedule s = out;
      array<bool, D> moved;
      fill(ALL(moved), false);
      int cur = j;
      moved.at(cur) = true;
      while (1) {
        int max_j2 = -1;
        int max_g = numeric_limits<int>::min();
        for (int j2 = max(0, cur - 10); j2 < min(D, cur + 10); ++j2) {
          if (moved.at(j2) || s.get(j2) == s.get(cur)) {
            continue;
          }
          int g = calc_gain_swap(s, j, j2);
          if (max_g < g) {
            max_j2 = j2;
            max_g = g;
          }
        }
        if (max_j2 == -1) {
          break;
        }
        cur = max_j2;
        s.swap(j, cur);
        moved.at(cur) = true;
        if (s.score > out.score) {
          out = s;
        }
        if (get_cur_sec() >= exit_time) {
          return;
        }
      }
    }
  }
}

void search_sa(Schedule &out, double exit_time, double temp_hi = 2000,
               double temp_lo = 600) {
  Schedule ans = out;
  uniform_real_distribution<> ur(0.0, 1.0);
  const double start_sec = get_cur_sec();
  double temp = temp_hi;
  for (int it = 0;; ++it) {
    shuffle(ALL(day_order), rng);
    // shuffle(ALL(type_order), rng);
    shuffle(ALL(type_order2), rng);

    for (int j : day_order) {
      // for(int k : type_order){
      for (int ki : type_order2) {
        const int k = cand_types.at(j).at(ki);
        if (k == ans.get(j)) {
          continue;
        }
        const int g = ans.calc_gain_set(j, k);
        if (exp((double)g / temp) >= ur(rng)) {
          ans.set(j, k);
          if (ans.score > out.score) {
            out = ans;
          }
        }
      }

      for (int delta = 1; delta <= MAX_DELTA_DAY && j + delta < D; ++delta) {
        const int t1 = ans.get(j), t2 = ans.get(j + delta);
        if (!(cand_types_set.at(j).count(t2) &&
              cand_types_set.at(j + delta).count(t1))) {
          continue;
        }
        const int g = calc_gain_swap(ans, j, j + delta);
        if (exp((double)g / temp) >= ur(rng)) {
          ans.swap(j, j + delta);
          if (ans.score > out.score) {
            out = ans;
          }
        }
      }
    }

    const double t = (get_cur_sec() - start_sec) / (exit_time - start_sec);
    temp = pow(temp_hi, 1 - t) * pow(temp_lo, t);
    if (t >= 1.0) {
      // cerr << it << endl;
      break;
    }
  }
}

void solve() {
  Schedule ans;
  greedy_init(ans, 1);
  // greedy_update(ans);
  search_sa(ans, 1.9);
  // search_cycle(ans, 1.95, 1);
  ans.output();
}

int main() {
  start = clock();
  rng.seed(0);
  input();
  init();
  solve();
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using ull = unsigned long long;
#define REP(i, n) for (int i = 0; i < (n); ++i)
#define RREP(i, n) for (int i = (n)-1; i >= 0; --i)
#define ALL(v) (v).begin(), (v).end()
template <class T>
using reverse_priority_queue = priority_queue<T, vector<T>, greater<T>>;

const int D = 365;
const int NTYPE = 26;
const int MAX_DELTA_DAY = 10;

minstd_rand rng;
clock_t start;

array<int, NTYPE> c;
array<array<int, NTYPE>, D> s;

const int NUM_CAND_TYPE = NTYPE / 2;
array<array<int, NUM_CAND_TYPE>, D> cand_types;
array<set<int>, D> cand_types_set;

array<int, D> day_order;
array<int, NTYPE> type_order;
array<int, NUM_CAND_TYPE> type_order2;

struct TypeSchedule {
  int id;
  set<int> day;
  int score;

  int get_prev_day(int x) const {
    auto it1 = day.lower_bound(x);
    int prev_day = -1;
    if (it1 != day.begin()) {
      prev_day = *(--it1);
    }
    return prev_day;
  }

  int get_next_day(int x) const {
    auto it2 = day.upper_bound(x);
    if (it2 != day.end() && *it2 == x) {
      ++it2;
    }
    int next_day = D;
    if (it2 != day.end()) {
      next_day = *it2;
    }
    return next_day;
  }

  TypeSchedule(int id) : id(id) { score = -c.at(id) * (D + 1) * D / 2; }

  int add_gain(int x) const {
    const int prev_day = get_prev_day(x);
    const int next_day = get_next_day(x);
    // prev_day + 1, ... , next_day - 1 = next_day - prev_day - 1 個
    const int l = next_day - prev_day - 1;
    int g = c.at(id) * (l + 1) * l / 2;
    // prev_day + 1, ..., x - 1 = x - prev_day - 1 個
    const int m1 = x - prev_day - 1;
    const int m2 = next_day - x - 1;
    g -= c.at(id) * (m1 * (m1 + 1) / 2 + m2 * (m2 + 1) / 2);
    g += s.at(x).at(id);
    return g;
  }

  void add(int x) {
    score += add_gain(x);
    day.insert(x);
  }

  int remove_gain(int x) const {
    assert(day.count(x) == 1);
    const int prev_day = get_prev_day(x);
    const int next_day = get_next_day(x);
    // prev_day + 1, ... , next_day - 1 = next_day - prev_day - 1 個
    const int l = next_day - prev_day - 1;
    int g = -c.at(id) * (l + 1) * l / 2;
    // prev_day + 1, ..., x - 1 = x - prev_day - 1 個
    const int m1 = x - prev_day - 1;
    const int m2 = next_day - x - 1;
    g += c.at(id) * (m1 * (m1 + 1) / 2 + m2 * (m2 + 1) / 2);
    g -= s.at(x).at(id);
    return g;
  }

  void remove(int x) {
    score += remove_gain(x);
    day.erase(x);
  }
};

struct Schedule {
  vector<TypeSchedule> type;
  array<int, D> d2t;
  int score;
  Schedule() {
    fill(ALL(d2t), -1);
    score = 0;
    REP(i, NTYPE) {
      type.emplace_back(i);
      score += type.at(i).score;
    }
    validate_score();
  }

  int calc_gain_set(int d, int t) const {
    const int t0 = d2t.at(d);
    if (t0 == t) {
      return 0;
    }
    int ret = 0;
    if (t0 != -1) {
      ret += type.at(t0).remove_gain(d);
    }
    ret += type.at(t).add_gain(d);
    return ret;
  }

  void set(int d, int t) {
    if (d2t.at(d) == t) {
      return;
    }
    score += calc_gain_set(d, t);
    const int t0 = d2t.at(d);
    d2t.at(d) = t;
    type.at(t).add(d);
    if (t0 != -1) {
      type.at(t0).remove(d);
    }
    validate_score();
  }

  int get(int d) const { return d2t.at(d); }
  void unset(int d, int t) {
    score -= type.at(t).score;
    d2t.at(d) = -1;
    type.at(t).remove(d);
    score += type.at(t).score;
    validate_score();
  }
  void swap(int d1, int d2) {
    if (get(d1) == get(d2)) {
      return;
    }
    const int t1 = get(d1);
    const int t2 = get(d2);
    set(d1, t2);
    set(d2, t1);
    validate_score();
  }

  void output() const {
    REP(i, D) {
      assert(d2t.at(i) != -1);
      cout << d2t.at(i) + 1 << endl;
    }
  }

  void validate_score() const {
    return;
    // int tmp = 0;
    // array<int, NTYPE> last;
    // fill(ALL(last), 0);
    // REP(i, D){
    //     int t = get(i);
    //     if(t != -1){
    //         tmp += s.at(i).at(t);
    //         last.at(t) = i + 1;
    //     }
    //     REP(s, NTYPE){
    //         tmp -= c.at(s) * (i + 1 - last.at(s));
    //     }
    // }
    // assert(tmp == score);
  }
};

void input() {
  int _d;
  cin >> _d;
  REP(i, NTYPE) { cin >> c.at(i); }
  REP(i, D) {
    REP(j, NTYPE) { cin >> s.at(i).at(j); }
  }
}

int calc_score_type(const array<int, D> &sch, int t) {
  int last = 0;
  int ret = 0;
  REP(i, D) {
    if (sch.at(i) == t) {
      ret += s.at(i).at(t);
      last = i + 1;
    }
    ret -= c.at(t) * (i + 1 - last);
  }
  return ret;
}

int calc_satisfaction(const array<int, D> &sch) {
  int ret = 0;
  REP(i, NTYPE) { ret += calc_score_type(sch, i); }
  return ret;
}

void greedy_init(Schedule &ans, int k = 1) {
  array<int, NTYPE> last;
  fill(ALL(last), 0);
  REP(i, D) {
    int max_type = -1;
    int max_val = numeric_limits<int>::min();
    REP(j, NTYPE) {
      int tmp = s.at(i).at(j);
      tmp += c.at(j) * (i + 1 - last.at(j)) * (min(D, i + k + 1) - i);
      if (max_val < tmp) {
        max_type = j;
        max_val = tmp;
      }
    }
    assert(max_type != -1);
    ans.set(i, max_type);
    last.at(max_type) = i + 1;
  }
}

void init() {
  REP(i, D) { day_order.at(i) = i; }
  REP(i, NTYPE) { type_order.at(i) = i; }
  REP(i, NUM_CAND_TYPE) { type_order2.at(i) = i; }

  REP(i, D) {
    array<pair<int, int>, NTYPE> tmp;
    REP(j, NTYPE) { tmp.at(j) = {s.at(i).at(j) + c.at(j), j}; }
    sort(ALL(tmp), greater<pair<int, int>>());
    REP(j, NUM_CAND_TYPE) {
      cand_types.at(i).at(j) = tmp.at(j).second;
      cand_types_set.at(i).insert(tmp.at(j).second);
    }
  }
}

double get_cur_sec() { return (double)(clock() - start) / CLOCKS_PER_SEC; }

int calc_gain_swap(Schedule &ans, int i, int j) {
  if (ans.get(i) == ans.get(j)) {
    return 0;
  }
  const int ti = ans.get(i);
  const int tj = ans.get(j);
  assert(ti != -1 && tj != -1);
  int g = -ans.score;
  ans.set(i, tj);
  ans.set(j, ti);
  g += ans.score;
  ans.set(i, ti);
  ans.set(j, tj);
  return g;
}

void greedy_update(Schedule &ans) {
  bool modified = true;
  while (modified) {
    modified = false;
    // change type
    REP(j, D) {
      REP(k, NTYPE) {
        if (ans.get(j) == k) {
          continue;
        }
        const int g = ans.calc_gain_set(j, k);
        if (g > 0) {
          modified = true;
          ans.set(j, k);
        }
      }
    }

    // swap day
    REP(j, D) {
      for (int delta = 1; delta <= MAX_DELTA_DAY && j + delta < D; ++delta) {
        const int g = calc_gain_swap(ans, j, j + delta);
        if (g > 0) {
          modified = true;
          ans.swap(j, j + delta);
        }
      }
    }

    if (get_cur_sec() >= 1.9) {
      break;
    }
  }
}

void search_cycle(Schedule &out, double exit_time, int max_it) {
  for (int it = 0; it < max_it; ++it) {
    REP(j, D) {
      Schedule s = out;
      array<bool, D> moved;
      fill(ALL(moved), false);
      int cur = j;
      moved.at(cur) = true;
      while (1) {
        int max_j2 = -1;
        int max_g = numeric_limits<int>::min();
        for (int j2 = max(0, cur - 10); j2 < min(D, cur + 10); ++j2) {
          if (moved.at(j2) || s.get(j2) == s.get(cur)) {
            continue;
          }
          int g = calc_gain_swap(s, j, j2);
          if (max_g < g) {
            max_j2 = j2;
            max_g = g;
          }
        }
        if (max_j2 == -1) {
          break;
        }
        cur = max_j2;
        s.swap(j, cur);
        moved.at(cur) = true;
        if (s.score > out.score) {
          out = s;
        }
        if (get_cur_sec() >= exit_time) {
          return;
        }
      }
    }
  }
}

void search_sa(Schedule &out, double exit_time, double temp_hi = 2000,
               double temp_lo = 600) {
  Schedule ans = out;
  uniform_real_distribution<> ur(0.0, 1.0);
  const double start_sec = get_cur_sec();
  double temp = temp_hi;
  for (int it = 0;; ++it) {
    shuffle(ALL(day_order), rng);
    // shuffle(ALL(type_order), rng);
    shuffle(ALL(type_order2), rng);

    for (int j : day_order) {
      // for(int k : type_order){
      for (int ki : type_order2) {
        const int k = cand_types.at(j).at(ki);
        if (k == ans.get(j)) {
          continue;
        }
        const int g = ans.calc_gain_set(j, k);
        if (exp((double)g / temp) >= ur(rng)) {
          ans.set(j, k);
          if (ans.score > out.score) {
            out = ans;
          }
        }
      }

      for (int delta = 1; delta <= MAX_DELTA_DAY && j + delta < D; ++delta) {
        const int t1 = ans.get(j), t2 = ans.get(j + delta);
        if (!(cand_types_set.at(j).count(t2) &&
              cand_types_set.at(j + delta).count(t1))) {
          continue;
        }
        const int g = calc_gain_swap(ans, j, j + delta);
        if (exp((double)g / temp) >= ur(rng)) {
          ans.swap(j, j + delta);
          if (ans.score > out.score) {
            out = ans;
          }
        }
      }
    }

    const double t = (get_cur_sec() - start_sec) / (exit_time - start_sec);
    temp = pow(temp_hi, 1 - t) * pow(temp_lo, t);
    if (t >= exit_time) {
      // cerr << it << endl;
      break;
    }
  }
}

void solve() {
  Schedule ans;
  greedy_init(ans, 1);
  // greedy_update(ans);
  search_sa(ans, 1.9);
  // search_cycle(ans, 1.95, 1);
  ans.output();
}

int main() {
  start = clock();
  rng.seed(0);
  input();
  init();
  solve();
  return 0;
}