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

#include <algorithm>
#include <bits/stdc++.h>
#include <cctype>
#include <cmath>
#include <functional>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <stdint.h>
#include <string>
#include <tuple>
#include <unordered_map>
#include <vector>

using namespace std;

typedef long double ld;
typedef int_fast64_t ll;
typedef pair<ll, ll> pii;
typedef pair<ll, pii> piii;
typedef pair<pii, pii> piiii;
typedef pair<ll, bool> pib;

typedef vector<bool> vb;

typedef vector<ll> vi;
typedef vector<vi> vvi;
typedef vector<vvi> vvvi;
typedef vector<pii> vpii;
typedef vector<vpii> vvpii;

typedef vector<double> vf;

typedef queue<ll> qi;

typedef long double ld;

#define mp make_pair
#define mt make_tuple
#define pb push_back

#define rep(i, N) for (ll i = 0; i < N; i++)
#define repa(i, a, N) for (ll i = a; i < N; i++)
#define repr(i, N) for (ll i = N - 1; i >= 0; i--)

#define drep(i, N, j, M) rep(i, N) rep(j, M)
#define trep(i, N, j, M, k, L) rep(i, N) rep(j, M) rep(k, L)

#define all(A) A.begin(), A.end()
#define uni(A) A.erase(unique(all(A)), A.end())

const ll LLMAX = numeric_limits<ll>::max() / 3;
const ld LDMAX = numeric_limits<ld>::max() / 3;
const ld PI = 3.141592653589793;
const ll MOD = 1e9 + 7;
// const ll MOD = 998244353;

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 const ll &value() const noexcept { return a; }
  constexpr modint operator+() const noexcept { return *this; }
  constexpr modint operator-() const noexcept {
    return modint(0LL) -= modint(*this);
  }
  constexpr modint operator+(const modint rhs) const noexcept {
    return modint(*this) += rhs;
  }
  constexpr modint operator-(const modint rhs) const noexcept {
    return modint(*this) -= rhs;
  }
  constexpr modint operator*(const modint rhs) const noexcept {
    return modint(*this) *= rhs;
  }
  constexpr modint operator/(const modint 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;
  }
};
typedef modint<MOD> mod;

constexpr const mod operator+(const ll a, const mod b) noexcept {
  return mod(a) += b;
}
constexpr const mod operator-(const ll a, const mod b) noexcept {
  return mod(a) -= b;
}
constexpr const mod operator*(const ll a, const mod b) noexcept {
  return mod(a) *= b;
}
constexpr const mod operator/(const ll a, const mod b) noexcept {
  return mod(a) /= b;
}

template <class T> class SegTree {
  ll n;
  vector<T> data;
  T def;
  function<T(T, T)> operation;
  function<T(T, T)> update;
  T _query(ll a, ll b, ll k, ll l, ll r) {
    if (r <= a || b <= l)
      return def;
    if (a <= l && r <= b)
      return data[k];
    else {
      T c1 = _query(a, b, 2 * k + 1, l, (l + r) / 2);
      T c2 = _query(a, b, 2 * k + 2, (l + r) / 2, r);
      return operation(c1, c2);
    }
  }

public:
  // _n: サイズ, _def: 初期値・単位元, _operation: クエリ, _update: 更新
  SegTree(ll _n, T _def, function<T(T, T)> _operation,
          function<T(T, T)> _update)
      : def(_def), operation(_operation), update(_update) {
    n = 1;
    while (n < _n) {
      n *= 2;
    }
    data = vector<T>(2 * n - 1, def);
  }
  // iの値をxに変更
  void change(ll i, T x) {
    i += n - 1;
    data[i] = update(data[i], x);
    while (i > 0) {
      i = (i - 1) / 2;
      data[i] = operation(data[i * 2 + 1], data[i * 2 + 2]);
    }
  }
  // [a, b)の区間クエリ実行
  T query(ll a, ll b) { return _query(a, b, 0, 0, n); }
  // 添え字アクセス
  constexpr const T operator[](ll i) const noexcept { return data[i + n - 1]; }
  constexpr const ll size() const noexcept { return n; }
  constexpr const T unit() const noexcept { return def; }
};

/*
SegTree<ll> st(N, LLMAX,
        // a: left child, b: right child
        [](ll a, ll b) {return min(a, b); },
        // d: old, x: query
        [](ll d, ll x) {return x; });
// */

template <class t, class u>
ostream &operator<<(ostream &os, const pair<t, u> &p) {
  return os << "{" << p.first << "," << p.second << "}";
}
template <class t> ostream &operator<<(ostream &os, const vector<t> &v) {
  os << "{";
  for (auto e : v)
    os << e << ",";
  return os << "}" << endl;
}
template <class t, class u>
ostream &operator<<(ostream &os, const unordered_map<t, u> &m) {
  os << "{";
  for (auto itr = m.begin(); itr != m.end(); ++itr) {
    os << "(" << itr->first << ", " << itr->second << "), ";
  }
  return os << "}";
}
template <class t> ostream &operator<<(ostream &os, const set<t> &m) {
  os << "{";
  for (auto itr = m.begin(); itr != m.end(); ++itr) {
    os << *itr << ", ";
  }
  return os << "}";
}
template <class t> ostream &operator<<(ostream &os, const multiset<t> &m) {
  os << "{";
  for (auto itr = m.begin(); itr != m.end(); ++itr) {
    os << *itr << ", ";
  }
  return os << "}";
}
template <ll Modulus>
ostream &operator<<(ostream &os, const modint<Modulus> &m) {
  return os << (m.a + Modulus) % Modulus;
}

template <class t> ostream &operator<<(ostream &os, const SegTree<t> &st) {
  os << "{";
  ll s = st.size();
  rep(i, s) {
    t v = st[i];
    if (v == st.unit()) {
      os << "-, ";
    } else {
      os << v << ", ";
    }
  }
  return os << "}\\n";
}

typedef vector<mod> vm;
typedef vector<vm> vvm;
typedef vector<vvm> vvvm;

template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}

void No() { cout << "No" << endl; }
void Yes() { cout << "Yes" << endl; }
void NO() { cout << "NO" << endl; }
void YES() { cout << "YES" << endl; }

// x^n(mod M) ←普通にpow(x,n)では溢れてしまうため，随時mod計算
ll mpow(ll x, ll n) {
  ll ans = 1;
  while (n != 0) {
    if (n & 1)
      ans = ans * x % MOD;
    x = x * x % MOD;
    n = n >> 1;
  }
  return ans;
}

ll fac_size = 1e6;
vm fac(fac_size + 1);  // n! (mod MOD)
vm ifac(fac_size + 1); // k!^{M-2} (mod MOD)
bool fac_calculated = false;

void pre_calc_comb() {
  fac[0] = 1;
  ifac[0] = 1;
  for (ll i = 0; i < fac_size; i++) {
    fac[i + 1] = fac[i] * (i + 1); // n!(mod M)
    ifac[i + 1] =
        ifac[i] * mpow(i + 1, MOD - 2); // k!^{M-2} (mod M) ←累乗にmpowを採用
  }
  fac_calculated = true;
}

mod comb(ll a, ll b) { // aCbをmod計算
  if (!fac_calculated)
    pre_calc_comb();

  if (a == 0 && b == 0)
    return 1;
  if (a < b || a < 0)
    return 0;
  mod tmp = ifac[a - b] * ifac[b];
  return tmp * fac[a];
}

ll gcd(ll a, ll b) {
  if (a == 0)
    return b;
  if (b == 0)
    return a;

  if (a % b == 0) {
    return (b);
  } else {
    return (gcd(b, a % b));
  }
}

ll N, M, K, Q;
ll ans = 0;
ll H, W;
string S;

void solve() {
  ll D;
  cin >> D;

  vi C(26);
  rep(i, 26) { cin >> C[i]; }

  vvi S(D, vi(26));
  rep(d, D) {
    rep(i, 26) { cin >> S[d][i]; }
  }

  vi T(D, 0);

  auto calc = [&](ll D) {
    ll ans = 0;
    vi L(26, -1);
    // 貪欲
    rep(d, D) {
      ll t = T[d];
      L[t] = d;
      ans += S[d][t];

      rep(i, 26) { ans -= (d - L[i]) * C[i]; }
    }
    return ans;
  };

  // 初期値
  rep(d, D) {
    ll m = 0;
    ll mi = 0;
    rep(i, 26) {
      T[d] = i;
      ll tmp = calc(d + 1);
      if (m < tmp) {
        mi = i;
        m = tmp;
      }
    }
    T[d] = mi;
  }

  // シミュレーション
  vector<set<ll>> Vs(26);
  rep(i, 26) {
    Vs[i].insert(-1);
    Vs[i].insert(D);
  }
  rep(d, D) { Vs[T[d]].insert(d); }

  auto diff = [&](ll d, ll next) {
    ll current = T[d];

    if (current == next) {
      ll zero = 0;
      return zero;
    }

    set<ll> &pre = Vs[current];
    set<ll> &post = Vs[next];

    auto pr_itr = pre.upper_bound(d);
    auto post_itr = post.upper_bound(d);
    ll pr = *pr_itr;
    ll po = *post_itr;

    pr_itr--;
    pr_itr--;
    post_itr--;
    ll nr = *pr_itr;
    ll no = *post_itr;

    ll di = (pr - d) * C[current] * (d - nr) - (po - d) * C[next] * (d - no);
    di += S[d][current] - S[d][next];
    return -di;
  };

  auto update = [&](ll d, ll next) {
    ll current = T[d];
    if (current == next)
      return;

    set<ll> &pre = Vs[current];
    set<ll> &post = Vs[next];
    pre.erase(d);
    post.insert(d);
    T[d] = next;
  };

  ll M = 8e5;
  vpii DQ(M);
  random_device rnd;
  rep(i, M) {
    DQ[i].first = rnd() % D;
    DQ[i].second = rnd() % 26;
  }

  ans = calc(D);
  rep(m, M) {
    ll d = DQ[m].first;
    ll md = 0;
    ll mi = T[d];

    rep(i, 26) {
      ll di = diff(d, i);
      if (md < di) {
        mi = i;
        md = di;
      }
    }

    // 改善
    ans += md;
    update(d, mi);
  }

  rep(d, D) { cout << T[d] + 1 << endl; }
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  solve();
  return 0;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <bits/stdc++.h>
#include <cctype>
#include <cmath>
#include <functional>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <stdint.h>
#include <string>
#include <tuple>
#include <unordered_map>
#include <vector>

using namespace std;

typedef long double ld;
typedef int_fast64_t ll;
typedef pair<ll, ll> pii;
typedef pair<ll, pii> piii;
typedef pair<pii, pii> piiii;
typedef pair<ll, bool> pib;

typedef vector<bool> vb;

typedef vector<ll> vi;
typedef vector<vi> vvi;
typedef vector<vvi> vvvi;
typedef vector<pii> vpii;
typedef vector<vpii> vvpii;

typedef vector<double> vf;

typedef queue<ll> qi;

typedef long double ld;

#define mp make_pair
#define mt make_tuple
#define pb push_back

#define rep(i, N) for (ll i = 0; i < N; i++)
#define repa(i, a, N) for (ll i = a; i < N; i++)
#define repr(i, N) for (ll i = N - 1; i >= 0; i--)

#define drep(i, N, j, M) rep(i, N) rep(j, M)
#define trep(i, N, j, M, k, L) rep(i, N) rep(j, M) rep(k, L)

#define all(A) A.begin(), A.end()
#define uni(A) A.erase(unique(all(A)), A.end())

const ll LLMAX = numeric_limits<ll>::max() / 3;
const ld LDMAX = numeric_limits<ld>::max() / 3;
const ld PI = 3.141592653589793;
const ll MOD = 1e9 + 7;
// const ll MOD = 998244353;

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 const ll &value() const noexcept { return a; }
  constexpr modint operator+() const noexcept { return *this; }
  constexpr modint operator-() const noexcept {
    return modint(0LL) -= modint(*this);
  }
  constexpr modint operator+(const modint rhs) const noexcept {
    return modint(*this) += rhs;
  }
  constexpr modint operator-(const modint rhs) const noexcept {
    return modint(*this) -= rhs;
  }
  constexpr modint operator*(const modint rhs) const noexcept {
    return modint(*this) *= rhs;
  }
  constexpr modint operator/(const modint 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;
  }
};
typedef modint<MOD> mod;

constexpr const mod operator+(const ll a, const mod b) noexcept {
  return mod(a) += b;
}
constexpr const mod operator-(const ll a, const mod b) noexcept {
  return mod(a) -= b;
}
constexpr const mod operator*(const ll a, const mod b) noexcept {
  return mod(a) *= b;
}
constexpr const mod operator/(const ll a, const mod b) noexcept {
  return mod(a) /= b;
}

template <class T> class SegTree {
  ll n;
  vector<T> data;
  T def;
  function<T(T, T)> operation;
  function<T(T, T)> update;
  T _query(ll a, ll b, ll k, ll l, ll r) {
    if (r <= a || b <= l)
      return def;
    if (a <= l && r <= b)
      return data[k];
    else {
      T c1 = _query(a, b, 2 * k + 1, l, (l + r) / 2);
      T c2 = _query(a, b, 2 * k + 2, (l + r) / 2, r);
      return operation(c1, c2);
    }
  }

public:
  // _n: サイズ, _def: 初期値・単位元, _operation: クエリ, _update: 更新
  SegTree(ll _n, T _def, function<T(T, T)> _operation,
          function<T(T, T)> _update)
      : def(_def), operation(_operation), update(_update) {
    n = 1;
    while (n < _n) {
      n *= 2;
    }
    data = vector<T>(2 * n - 1, def);
  }
  // iの値をxに変更
  void change(ll i, T x) {
    i += n - 1;
    data[i] = update(data[i], x);
    while (i > 0) {
      i = (i - 1) / 2;
      data[i] = operation(data[i * 2 + 1], data[i * 2 + 2]);
    }
  }
  // [a, b)の区間クエリ実行
  T query(ll a, ll b) { return _query(a, b, 0, 0, n); }
  // 添え字アクセス
  constexpr const T operator[](ll i) const noexcept { return data[i + n - 1]; }
  constexpr const ll size() const noexcept { return n; }
  constexpr const T unit() const noexcept { return def; }
};

/*
SegTree<ll> st(N, LLMAX,
        // a: left child, b: right child
        [](ll a, ll b) {return min(a, b); },
        // d: old, x: query
        [](ll d, ll x) {return x; });
// */

template <class t, class u>
ostream &operator<<(ostream &os, const pair<t, u> &p) {
  return os << "{" << p.first << "," << p.second << "}";
}
template <class t> ostream &operator<<(ostream &os, const vector<t> &v) {
  os << "{";
  for (auto e : v)
    os << e << ",";
  return os << "}" << endl;
}
template <class t, class u>
ostream &operator<<(ostream &os, const unordered_map<t, u> &m) {
  os << "{";
  for (auto itr = m.begin(); itr != m.end(); ++itr) {
    os << "(" << itr->first << ", " << itr->second << "), ";
  }
  return os << "}";
}
template <class t> ostream &operator<<(ostream &os, const set<t> &m) {
  os << "{";
  for (auto itr = m.begin(); itr != m.end(); ++itr) {
    os << *itr << ", ";
  }
  return os << "}";
}
template <class t> ostream &operator<<(ostream &os, const multiset<t> &m) {
  os << "{";
  for (auto itr = m.begin(); itr != m.end(); ++itr) {
    os << *itr << ", ";
  }
  return os << "}";
}
template <ll Modulus>
ostream &operator<<(ostream &os, const modint<Modulus> &m) {
  return os << (m.a + Modulus) % Modulus;
}

template <class t> ostream &operator<<(ostream &os, const SegTree<t> &st) {
  os << "{";
  ll s = st.size();
  rep(i, s) {
    t v = st[i];
    if (v == st.unit()) {
      os << "-, ";
    } else {
      os << v << ", ";
    }
  }
  return os << "}\\n";
}

typedef vector<mod> vm;
typedef vector<vm> vvm;
typedef vector<vvm> vvvm;

template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}

void No() { cout << "No" << endl; }
void Yes() { cout << "Yes" << endl; }
void NO() { cout << "NO" << endl; }
void YES() { cout << "YES" << endl; }

// x^n(mod M) ←普通にpow(x,n)では溢れてしまうため，随時mod計算
ll mpow(ll x, ll n) {
  ll ans = 1;
  while (n != 0) {
    if (n & 1)
      ans = ans * x % MOD;
    x = x * x % MOD;
    n = n >> 1;
  }
  return ans;
}

ll fac_size = 1e6;
vm fac(fac_size + 1);  // n! (mod MOD)
vm ifac(fac_size + 1); // k!^{M-2} (mod MOD)
bool fac_calculated = false;

void pre_calc_comb() {
  fac[0] = 1;
  ifac[0] = 1;
  for (ll i = 0; i < fac_size; i++) {
    fac[i + 1] = fac[i] * (i + 1); // n!(mod M)
    ifac[i + 1] =
        ifac[i] * mpow(i + 1, MOD - 2); // k!^{M-2} (mod M) ←累乗にmpowを採用
  }
  fac_calculated = true;
}

mod comb(ll a, ll b) { // aCbをmod計算
  if (!fac_calculated)
    pre_calc_comb();

  if (a == 0 && b == 0)
    return 1;
  if (a < b || a < 0)
    return 0;
  mod tmp = ifac[a - b] * ifac[b];
  return tmp * fac[a];
}

ll gcd(ll a, ll b) {
  if (a == 0)
    return b;
  if (b == 0)
    return a;

  if (a % b == 0) {
    return (b);
  } else {
    return (gcd(b, a % b));
  }
}

ll N, M, K, Q;
ll ans = 0;
ll H, W;
string S;

void solve() {
  ll D;
  cin >> D;

  vi C(26);
  rep(i, 26) { cin >> C[i]; }

  vvi S(D, vi(26));
  rep(d, D) {
    rep(i, 26) { cin >> S[d][i]; }
  }

  vi T(D, 0);

  auto calc = [&](ll D) {
    ll ans = 0;
    vi L(26, -1);
    // 貪欲
    rep(d, D) {
      ll t = T[d];
      L[t] = d;
      ans += S[d][t];

      rep(i, 26) { ans -= (d - L[i]) * C[i]; }
    }
    return ans;
  };

  // 初期値
  rep(d, D) {
    ll m = 0;
    ll mi = 0;
    rep(i, 26) {
      T[d] = i;
      ll tmp = calc(d + 1);
      if (m < tmp) {
        mi = i;
        m = tmp;
      }
    }
    T[d] = mi;
  }

  // シミュレーション
  vector<set<ll>> Vs(26);
  rep(i, 26) {
    Vs[i].insert(-1);
    Vs[i].insert(D);
  }
  rep(d, D) { Vs[T[d]].insert(d); }

  auto diff = [&](ll d, ll next) {
    ll current = T[d];

    if (current == next) {
      ll zero = 0;
      return zero;
    }

    set<ll> &pre = Vs[current];
    set<ll> &post = Vs[next];

    auto pr_itr = pre.upper_bound(d);
    auto post_itr = post.upper_bound(d);
    ll pr = *pr_itr;
    ll po = *post_itr;

    pr_itr--;
    pr_itr--;
    post_itr--;
    ll nr = *pr_itr;
    ll no = *post_itr;

    ll di = (pr - d) * C[current] * (d - nr) - (po - d) * C[next] * (d - no);
    di += S[d][current] - S[d][next];
    return -di;
  };

  auto update = [&](ll d, ll next) {
    ll current = T[d];
    if (current == next)
      return;

    set<ll> &pre = Vs[current];
    set<ll> &post = Vs[next];
    pre.erase(d);
    post.insert(d);
    T[d] = next;
  };

  ll M = 1e6;
  vpii DQ(M);
  random_device rnd;
  rep(i, M) {
    DQ[i].first = rnd() % D;
    DQ[i].second = rnd() % 26;
  }

  ans = calc(D);
  rep(m, M) {
    ll d = DQ[m].first;
    ll md = 0;
    ll mi = T[d];

    rep(i, 26) {
      ll di = diff(d, i);
      if (md < di) {
        mi = i;
        md = di;
      }
    }

    // 改善
    ans += md;
    update(d, mi);
  }

  rep(d, D) { cout << T[d] + 1 << endl; }
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  solve();
  return 0;
}