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

#include <algorithm>
#include <cassert>
#include <chrono>
#include <cstdint>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <limits>
#include <map>
#include <math.h>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <unordered_map>
#include <vector>
#define REP(i, n) for (ll i = 0, i##_len = (n); i < i##_len; ++i)
#define RREP(i, n) for (ll i = (n)-1; i >= 0; --i)
#define SREP(i, s, n) for (ll i = (s), i##_len = (n); i < i##_len; ++i)
#define ALL(v) (v).begin(), (v).end()

using namespace std;
typedef long long int ll;
typedef unsigned long long int ull;

template <typename T> T gcd(T a, T b) {
  if (a < b)
    gcd(b, a);
  if (b == 1)
    return 1;
  T r;
  while ((r = a % b)) {
    a = b;
    b = r;
  }
  return b;
}

template <typename T> struct UnionFind {
  vector<T> par;
  vector<T> rank;

  UnionFind(T n = 1) { init(n); }

  void init(T n = 1) {
    par.resize(n);
    rank.resize(n);
    for (T i = 0; i < n; ++i) {
      par[i] = i;
      rank[i] = 0;
    }
  }

  int root(T x) {
    if (par[x] == x) {
      return x;
    } else {
      int r = root(par[x]);
      return par[x] = r;
    }
  }

  bool issame(T x, T y) { return root(x) == root(y); }

  bool merge(T x, T y) {
    x = root(x);
    y = root(y);
    if (x == y)
      return false;
    if (rank[x] < rank[y])
      swap(x, y);
    if (rank[x] == rank[y])
      ++rank[x];
    par[y] = x;
    return true;
  }
};

template <typename Monoid, typename OperatorMonoid = Monoid>
struct LazySegmentTree {
  using F = function<Monoid(Monoid, Monoid)>;
  using G = function<Monoid(Monoid, OperatorMonoid)>;
  using H = function<OperatorMonoid(OperatorMonoid, OperatorMonoid)>;

  ll sz, height;
  vector<Monoid> data;
  vector<OperatorMonoid> lazy;
  const F f;
  const G g;
  const H h;
  const Monoid M1;
  const OperatorMonoid OM0;

  LazySegmentTree(ll n, const F f, const G g, const H h, const Monoid &M1,
                  const OperatorMonoid OM0)
      : f(f), g(g), h(h), M1(M1), OM0(OM0) {
    sz = 1;
    height = 0;
    while (sz < n)
      sz <<= 1, height++;
    data.assign(2 * sz, M1);
    lazy.assign(2 * sz, OM0);
  }

  void reset() {
    data.assign(2 * sz, M1);
    lazy.assign(2 * sz, OM0);
  }

  void set(ll k, const Monoid &x) { data[k + sz] = x; }

  void build() {
    for (int k = sz - 1; k > 0; k--) {
      data[k] = f(data[2 * k + 0], data[2 * k + 1]);
    }
  }

  inline void propagate(ll k) {
    if (lazy[k] != OM0) {
      lazy[2 * k + 0] = h(lazy[2 * k + 0], lazy[k]);
      lazy[2 * k + 1] = h(lazy[2 * k + 1], lazy[k]);
      data[k] = reflect(k);
      lazy[k] = OM0;
    }
  }

  inline Monoid reflect(ll k) {
    return lazy[k] == OM0 ? data[k] : g(data[k], lazy[k]);
  }

  inline void recalc(ll k) {
    while (k >>= 1)
      data[k] = f(reflect(2 * k + 0), reflect(2 * k + 1));
  }

  inline void thrust(ll k) {
    for (int i = height; i > 0; i--)
      propagate(k >> i);
  }

  void update(ll a, ll b, const OperatorMonoid &x) {
    thrust(a += sz);
    thrust(b += sz - 1);
    for (int l = a, r = b + 1; l < r; l >>= 1, r >>= 1) {
      if (l & 1)
        lazy[l] = h(lazy[l], x), ++l;
      if (r & 1)
        --r, lazy[r] = h(lazy[r], x);
    }
    recalc(a);
    recalc(b);
  }

  Monoid query(ll a, ll b) {
    thrust(a += sz);
    thrust(b += sz - 1);
    Monoid L = M1, R = M1;
    for (int l = a, r = b + 1; l < r; l >>= 1, r >>= 1) {
      if (l & 1)
        L = f(L, reflect(l++));
      if (r & 1)
        R = f(reflect(--r), R);
    }
    return f(L, R);
  }

  Monoid operator[](const int &k) { return query(k, k + 1); }

  template <typename C>
  int find_subtree(ll a, const C &check, Monoid &M, bool type) {
    while (a < sz) {
      propagate(a);
      Monoid nxt =
          type ? f(reflect(2 * a + type), M) : f(M, reflect(2 * a + type));
      if (check(nxt))
        a = 2 * a + type;
      else
        M = nxt, a = 2 * a + 1 - type;
    }
    return a - sz;
  }

  template <typename C> int find_first(ll a, const C &check) {
    Monoid L = M1;
    if (a <= 0) {
      if (check(f(L, reflect(1))))
        return find_subtree(1, check, L, false);
      return -1;
    }
    thrust(a + sz);
    int b = sz;
    for (a += sz, b += sz; a < b; a >>= 1, b >>= 1) {
      if (a & 1) {
        Monoid nxt = f(L, reflect(a));
        if (check(nxt))
          return find_subtree(a, check, L, false);
        L = nxt;
        ++a;
      }
    }
    return -1;
  }

  template <typename C> int find_last(ll b, const C &check) {
    Monoid R = M1;
    if (b >= sz) {
      if (check(f(reflect(1), R)))
        return find_subtree(1, check, R, true);
      return -1;
    }
    thrust(b + sz - 1);
    int a = sz;
    for (b += sz; a < b; a >>= 1, b >>= 1) {
      if (b & 1) {
        Monoid nxt = f(reflect(--b), R);
        if (check(nxt))
          return find_subtree(b, check, R, true);
        R = nxt;
      }
    }
    return -1;
  }
};

int msb(ll x) {
  int n = 0;
  while (x > 0) {
    x /= 2;
    n++;
  }
  return n;
}

template <ll Modulus> class modint {
public:
  ll a;

  constexpr modint(const ll x = 0) noexcept : a(x % Modulus) {}
  constexpr ll &value() noexcept { return a; }
  constexpr operator int() const { return a; }
  constexpr operator long long() const { return a; }
  constexpr const ll &value() const noexcept { return a; }
  constexpr modint operator+(const modint rhs) const noexcept {
    return modint(*this) += rhs;
  }
  constexpr modint operator+(const ll rhs) const noexcept {
    return modint(*this) += rhs;
  }
  constexpr modint operator-(const modint rhs) const noexcept {
    return modint(*this) -= rhs;
  }
  constexpr modint operator-(const ll rhs) const noexcept {
    return modint(*this) -= rhs;
  }
  constexpr modint operator*(const modint rhs) const noexcept {
    return modint(*this) *= rhs;
  }
  constexpr modint operator*(const ll rhs) const noexcept {
    return modint(*this) *= rhs;
  }
  constexpr modint operator/(const modint rhs) const noexcept {
    return modint(*this) /= rhs;
  }
  constexpr modint operator/(const ll rhs) const noexcept {
    return modint(*this) /= rhs;
  }
  constexpr modint &operator+=(const modint rhs) noexcept {
    a += rhs.a;
    if (a >= Modulus) {
      a -= Modulus;
    }
    return *this;
  }
  constexpr modint &operator-=(const modint rhs) noexcept {
    if (a < rhs.a) {
      a += Modulus;
    }
    a -= rhs.a;
    return *this;
  }
  constexpr modint &operator*=(const modint rhs) noexcept {
    a = a * rhs.a % Modulus;
    return *this;
  }
  constexpr modint &operator/=(modint rhs) noexcept {
    ll exp = Modulus - 2;
    while (exp) {
      if (exp % 2) {
        *this *= rhs;
      }
      rhs *= rhs;
      exp /= 2;
    }
    return *this;
  }
};

vector<ll> read_vec(ll N, ll offset) {
  vector<ll> v(N);
  REP(i, N) {
    cin >> v[i];
    v[i] += offset;
  }
  return v;
}

#define MOD (ll(998244353))
#define INF (ll(1e15))

ll D;
vector<ll> c;
vector<vector<ll>> s;

ll score(vector<ll> &t, vector<ll>::iterator end) {
  vector<ll> last(26, -1);
  ll y = 0;
  auto it = t.begin();
  ll day = 0;
  while (it != end) {
    y += s[day][*it];
    last[*it] = day;
    REP(i, 26) { y -= c[i] * (day - last[i]); }
    it++;
    day++;
  }
  return y;
}

int main(int argc, const char *argv[]) {
  (void)argc;
  (void)argv;
  cin.tie(0);
  ios::sync_with_stdio(false);

  cin >> D;
  c = read_vec(26, 0);
  s.resize(D);
  REP(i, D) { s[i] = read_vec(26, 0); }

  auto t = vector<ll>(D);
  REP(i, D) { t[i] = max_element(ALL(s[i])) - s[i].begin(); }
  vector<set<ll>> ts(26);
  REP(i, 26) {
    ts[i].insert(-1);
    ts[i].insert(D);
  }
  REP(i, D) { ts[t[i]].insert(i); }

  ll x = score(t, t.end());
  ll M = 10000000;
  std::random_device seed_gen;
  std::mt19937 engine(seed_gen());

  REP(i, M) {
    ll d = engine() % D;
    ll q = t[d];
    while (q == t[d])
      q = engine() % 26;

    ll y = x;
    y -= s[d][t[d]];
    y += s[d][q];

    auto pit = ts[t[d]].find(d);
    ll ph = d - *prev(pit);
    ll pw = *next(pit) - d;

    y -= c[t[d]] * ph * pw;

    auto nit = prev(ts[q].upper_bound(d));
    ll nh = d - *nit;
    ll nw = *next(nit) - d;
    y += c[q] * nh * nw;

    if (y > x) {
      ts[t[d]].erase(pit);
      ts[q].insert(d);
      t[d] = q;
      x = y;
    }
  }

  //    cerr << x << endl;
  REP(i, D) { cout << t[i] + 1 << endl; }
}

// https://atcoder.jp/contests/intro-heuristics/tasks/intro_heuristics_a


Error: TLE

Fixed submission:

#include <algorithm>
#include <cassert>
#include <chrono>
#include <cstdint>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <limits>
#include <map>
#include <math.h>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <unordered_map>
#include <vector>
#define REP(i, n) for (ll i = 0, i##_len = (n); i < i##_len; ++i)
#define RREP(i, n) for (ll i = (n)-1; i >= 0; --i)
#define SREP(i, s, n) for (ll i = (s), i##_len = (n); i < i##_len; ++i)
#define ALL(v) (v).begin(), (v).end()

using namespace std;
typedef long long int ll;
typedef unsigned long long int ull;

template <typename T> T gcd(T a, T b) {
  if (a < b)
    gcd(b, a);
  if (b == 1)
    return 1;
  T r;
  while ((r = a % b)) {
    a = b;
    b = r;
  }
  return b;
}

template <typename T> struct UnionFind {
  vector<T> par;
  vector<T> rank;

  UnionFind(T n = 1) { init(n); }

  void init(T n = 1) {
    par.resize(n);
    rank.resize(n);
    for (T i = 0; i < n; ++i) {
      par[i] = i;
      rank[i] = 0;
    }
  }

  int root(T x) {
    if (par[x] == x) {
      return x;
    } else {
      int r = root(par[x]);
      return par[x] = r;
    }
  }

  bool issame(T x, T y) { return root(x) == root(y); }

  bool merge(T x, T y) {
    x = root(x);
    y = root(y);
    if (x == y)
      return false;
    if (rank[x] < rank[y])
      swap(x, y);
    if (rank[x] == rank[y])
      ++rank[x];
    par[y] = x;
    return true;
  }
};

template <typename Monoid, typename OperatorMonoid = Monoid>
struct LazySegmentTree {
  using F = function<Monoid(Monoid, Monoid)>;
  using G = function<Monoid(Monoid, OperatorMonoid)>;
  using H = function<OperatorMonoid(OperatorMonoid, OperatorMonoid)>;

  ll sz, height;
  vector<Monoid> data;
  vector<OperatorMonoid> lazy;
  const F f;
  const G g;
  const H h;
  const Monoid M1;
  const OperatorMonoid OM0;

  LazySegmentTree(ll n, const F f, const G g, const H h, const Monoid &M1,
                  const OperatorMonoid OM0)
      : f(f), g(g), h(h), M1(M1), OM0(OM0) {
    sz = 1;
    height = 0;
    while (sz < n)
      sz <<= 1, height++;
    data.assign(2 * sz, M1);
    lazy.assign(2 * sz, OM0);
  }

  void reset() {
    data.assign(2 * sz, M1);
    lazy.assign(2 * sz, OM0);
  }

  void set(ll k, const Monoid &x) { data[k + sz] = x; }

  void build() {
    for (int k = sz - 1; k > 0; k--) {
      data[k] = f(data[2 * k + 0], data[2 * k + 1]);
    }
  }

  inline void propagate(ll k) {
    if (lazy[k] != OM0) {
      lazy[2 * k + 0] = h(lazy[2 * k + 0], lazy[k]);
      lazy[2 * k + 1] = h(lazy[2 * k + 1], lazy[k]);
      data[k] = reflect(k);
      lazy[k] = OM0;
    }
  }

  inline Monoid reflect(ll k) {
    return lazy[k] == OM0 ? data[k] : g(data[k], lazy[k]);
  }

  inline void recalc(ll k) {
    while (k >>= 1)
      data[k] = f(reflect(2 * k + 0), reflect(2 * k + 1));
  }

  inline void thrust(ll k) {
    for (int i = height; i > 0; i--)
      propagate(k >> i);
  }

  void update(ll a, ll b, const OperatorMonoid &x) {
    thrust(a += sz);
    thrust(b += sz - 1);
    for (int l = a, r = b + 1; l < r; l >>= 1, r >>= 1) {
      if (l & 1)
        lazy[l] = h(lazy[l], x), ++l;
      if (r & 1)
        --r, lazy[r] = h(lazy[r], x);
    }
    recalc(a);
    recalc(b);
  }

  Monoid query(ll a, ll b) {
    thrust(a += sz);
    thrust(b += sz - 1);
    Monoid L = M1, R = M1;
    for (int l = a, r = b + 1; l < r; l >>= 1, r >>= 1) {
      if (l & 1)
        L = f(L, reflect(l++));
      if (r & 1)
        R = f(reflect(--r), R);
    }
    return f(L, R);
  }

  Monoid operator[](const int &k) { return query(k, k + 1); }

  template <typename C>
  int find_subtree(ll a, const C &check, Monoid &M, bool type) {
    while (a < sz) {
      propagate(a);
      Monoid nxt =
          type ? f(reflect(2 * a + type), M) : f(M, reflect(2 * a + type));
      if (check(nxt))
        a = 2 * a + type;
      else
        M = nxt, a = 2 * a + 1 - type;
    }
    return a - sz;
  }

  template <typename C> int find_first(ll a, const C &check) {
    Monoid L = M1;
    if (a <= 0) {
      if (check(f(L, reflect(1))))
        return find_subtree(1, check, L, false);
      return -1;
    }
    thrust(a + sz);
    int b = sz;
    for (a += sz, b += sz; a < b; a >>= 1, b >>= 1) {
      if (a & 1) {
        Monoid nxt = f(L, reflect(a));
        if (check(nxt))
          return find_subtree(a, check, L, false);
        L = nxt;
        ++a;
      }
    }
    return -1;
  }

  template <typename C> int find_last(ll b, const C &check) {
    Monoid R = M1;
    if (b >= sz) {
      if (check(f(reflect(1), R)))
        return find_subtree(1, check, R, true);
      return -1;
    }
    thrust(b + sz - 1);
    int a = sz;
    for (b += sz; a < b; a >>= 1, b >>= 1) {
      if (b & 1) {
        Monoid nxt = f(reflect(--b), R);
        if (check(nxt))
          return find_subtree(b, check, R, true);
        R = nxt;
      }
    }
    return -1;
  }
};

int msb(ll x) {
  int n = 0;
  while (x > 0) {
    x /= 2;
    n++;
  }
  return n;
}

template <ll Modulus> class modint {
public:
  ll a;

  constexpr modint(const ll x = 0) noexcept : a(x % Modulus) {}
  constexpr ll &value() noexcept { return a; }
  constexpr operator int() const { return a; }
  constexpr operator long long() const { return a; }
  constexpr const ll &value() const noexcept { return a; }
  constexpr modint operator+(const modint rhs) const noexcept {
    return modint(*this) += rhs;
  }
  constexpr modint operator+(const ll rhs) const noexcept {
    return modint(*this) += rhs;
  }
  constexpr modint operator-(const modint rhs) const noexcept {
    return modint(*this) -= rhs;
  }
  constexpr modint operator-(const ll rhs) const noexcept {
    return modint(*this) -= rhs;
  }
  constexpr modint operator*(const modint rhs) const noexcept {
    return modint(*this) *= rhs;
  }
  constexpr modint operator*(const ll rhs) const noexcept {
    return modint(*this) *= rhs;
  }
  constexpr modint operator/(const modint rhs) const noexcept {
    return modint(*this) /= rhs;
  }
  constexpr modint operator/(const ll rhs) const noexcept {
    return modint(*this) /= rhs;
  }
  constexpr modint &operator+=(const modint rhs) noexcept {
    a += rhs.a;
    if (a >= Modulus) {
      a -= Modulus;
    }
    return *this;
  }
  constexpr modint &operator-=(const modint rhs) noexcept {
    if (a < rhs.a) {
      a += Modulus;
    }
    a -= rhs.a;
    return *this;
  }
  constexpr modint &operator*=(const modint rhs) noexcept {
    a = a * rhs.a % Modulus;
    return *this;
  }
  constexpr modint &operator/=(modint rhs) noexcept {
    ll exp = Modulus - 2;
    while (exp) {
      if (exp % 2) {
        *this *= rhs;
      }
      rhs *= rhs;
      exp /= 2;
    }
    return *this;
  }
};

vector<ll> read_vec(ll N, ll offset) {
  vector<ll> v(N);
  REP(i, N) {
    cin >> v[i];
    v[i] += offset;
  }
  return v;
}

#define MOD (ll(998244353))
#define INF (ll(1e15))

ll D;
vector<ll> c;
vector<vector<ll>> s;

ll score(vector<ll> &t, vector<ll>::iterator end) {
  vector<ll> last(26, -1);
  ll y = 0;
  auto it = t.begin();
  ll day = 0;
  while (it != end) {
    y += s[day][*it];
    last[*it] = day;
    REP(i, 26) { y -= c[i] * (day - last[i]); }
    it++;
    day++;
  }
  return y;
}

int main(int argc, const char *argv[]) {
  (void)argc;
  (void)argv;
  cin.tie(0);
  ios::sync_with_stdio(false);

  cin >> D;
  c = read_vec(26, 0);
  s.resize(D);
  REP(i, D) { s[i] = read_vec(26, 0); }

  auto t = vector<ll>(D);
  REP(i, D) { t[i] = max_element(ALL(s[i])) - s[i].begin(); }
  vector<set<ll>> ts(26);
  REP(i, 26) {
    ts[i].insert(-1);
    ts[i].insert(D);
  }
  REP(i, D) { ts[t[i]].insert(i); }

  ll x = score(t, t.end());
  ll M = 20000000;
  std::random_device seed_gen;
  std::mt19937 engine(seed_gen());

  REP(i, M) {
    ll d = engine() % D;
    ll q = t[d];
    while (q == t[d])
      q = engine() % 26;

    ll y = x;
    y -= s[d][t[d]];
    y += s[d][q];

    auto pit = ts[t[d]].find(d);
    ll ph = d - *prev(pit);
    ll pw = *next(pit) - d;

    y -= c[t[d]] * ph * pw;

    auto nit = prev(ts[q].upper_bound(d));
    ll nh = d - *nit;
    ll nw = *next(nit) - d;
    y += c[q] * nh * nw;

    if (y > x) {
      ts[t[d]].erase(pit);
      ts[q].insert(d);
      t[d] = q;
      x = y;
    }
  }

  //    cerr << x << endl;
  REP(i, D) { cout << t[i] + 1 << endl; }
}

// https://atcoder.jp/contests/intro-heuristics/tasks/intro_heuristics_a