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

#include <bits/stdc++.h>
using namespace std;

typedef unsigned long long ull;
typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
typedef pair<double, double> pdd;
// typedef vector<vector<ll>> Graph;

const ll mod = 1e9 + 7;
// const ll mod = 998244353;

#define REP(i, n) for (ll i = 0; i < (ll)n; i++)
#define dump(x) cerr << #x << " = " << (x) << endl;
#define spa << " " <<
#define fi first
#define se second

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

template <class S, class T>
ostream &operator<<(ostream &os, const pair<S, T> v) {
  os << "(" << v.first << ", " << v.second << ")";
  return os;
}
template <class T> ostream &operator<<(ostream &os, const vector<T> v) {
  for (int i = 0; i < (int)v.size(); i++) {
    if (i > 0) {
      os << " ";
    }
    os << v[i];
  }
  return os;
}
template <class T> ostream &operator<<(ostream &os, const vector<vector<T>> v) {
  for (int i = 0; i < (int)v.size(); i++) {
    if (i > 0) {
      os << endl;
    }
    os << v[i];
  }
  return os;
}

template <typename T> void debug(vector<vector<T>> &v, ll h, ll w) {
  for (ll i = 0; i < h; i++) {
    cerr << v[i][0];
    for (ll j = 1; j < w; j++)
      cerr spa v[i][j];
    cerr << endl;
  }
};
template <typename T> void debug(vector<T> &v, ll n) {
  if (n != 0)
    cerr << v[0];
  for (ll i = 1; i < n; i++)
    cerr spa v[i];
  cerr << endl;
};

string num2bit(ll num, ll len) {
  string bit = "";
  REP(i, len) { bit += char('0' + (num >> i & 1)); }
  return bit;
}

vector<ll> pros(vector<ll> &A) {
  ll N = A.size();
  vector<ll> B(N + 1, 0);
  REP(i, N) {
    ll lb = max(i - A[i], 0ll);
    ll ub = min(i + A[i], N - 1);
    B[lb]++;
    B[ub + 1]--;
  }
  REP(i, N) B[i + 1] += B[i];
  return B;
}

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

  ll N, K;
  cin >> N >> K;
  vector<ll> A(N);
  REP(i, N) cin >> A[i];

  REP(t, K) {
    auto B = pros(A);
    // cout << B << endl;
    REP(i, N) A[i] = B[i];

    bool ok = true;
    REP(i, N) {
      if (A[i] != N)
        ok = false;
    }
    if (ok)
      break;
  }

  REP(i, N) cout << A[i] << " ";
  cout << endl;

  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;

typedef unsigned long long ull;
typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
typedef pair<double, double> pdd;
// typedef vector<vector<ll>> Graph;

const ll mod = 1e9 + 7;
// const ll mod = 998244353;

#define REP(i, n) for (ll i = 0; i < (ll)n; i++)
#define dump(x) cerr << #x << " = " << (x) << endl;
#define spa << " " <<
#define fi first
#define se second

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

template <class S, class T>
ostream &operator<<(ostream &os, const pair<S, T> v) {
  os << "(" << v.first << ", " << v.second << ")";
  return os;
}
template <class T> ostream &operator<<(ostream &os, const vector<T> v) {
  for (int i = 0; i < (int)v.size(); i++) {
    if (i > 0) {
      os << " ";
    }
    os << v[i];
  }
  return os;
}
template <class T> ostream &operator<<(ostream &os, const vector<vector<T>> v) {
  for (int i = 0; i < (int)v.size(); i++) {
    if (i > 0) {
      os << endl;
    }
    os << v[i];
  }
  return os;
}

template <typename T> void debug(vector<vector<T>> &v, ll h, ll w) {
  for (ll i = 0; i < h; i++) {
    cerr << v[i][0];
    for (ll j = 1; j < w; j++)
      cerr spa v[i][j];
    cerr << endl;
  }
};
template <typename T> void debug(vector<T> &v, ll n) {
  if (n != 0)
    cerr << v[0];
  for (ll i = 1; i < n; i++)
    cerr spa v[i];
  cerr << endl;
};

string num2bit(ll num, ll len) {
  string bit = "";
  REP(i, len) { bit += char('0' + (num >> i & 1)); }
  return bit;
}

template <typename Monoid, typename OperatorMonoid, typename F, typename G,
          typename H>
struct LazySegmentTree {
  // std::function は遅いので使わないことを推奨
  // using F = function< Monoid(Monoid, Monoid) >;
  // using G = function< Monoid(Monoid, OperatorMonoid) >;
  // using H = function< OperatorMonoid(OperatorMonoid, OperatorMonoid) >;

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

  // n: サイズ
  // f: 二つの区間をマージする二項演算
  // g: 要素と作用素をマージする二項演算
  // h: 作用素同士をマージする二項演算
  // M1: モノイドの単位元
  // OM0: 作用素の単位元

  // 初期化を忘れない！x を初期値とすると
  // REP(i, n) seg[i].set(x);
  // seg.build();

  // 初期化なしだとM1で初期化される
  // range sum をするときは，M1は(0, 0)だが初期化は(0, 1)であることに注意！

  LazySegmentTree(int 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 set(int 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(int 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(int k) {
    return lazy[k] == OM0 ? data[k] : g(data[k], lazy[k]);
  }

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

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

  void update(int a, int 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(int a, int 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(int 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(int 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(int 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 main() {
  cin.tie(0);
  ios::sync_with_stdio(false);

  ll N, K;
  cin >> N >> K;
  vector<ll> A(N);
  REP(i, N) cin >> A[i];

  REP(t, K) {
    auto f = [](ll x, ll y) { return min(x, y); }; // なんでもいい
    auto g = [](ll x, ll a) { return x + a; };
    auto h = [](ll a, ll b) { return a + b; };
    ll M1 = (1ll << 60);
    ll OM0 = 0;
    LazySegmentTree<ll, ll, decltype(f), decltype(g), decltype(h)> seg(
        N, f, g, h, M1, OM0);
    // 値を初期化
    REP(i, N) seg.set(i, 0);
    seg.build();
    vector<ll> B(N, 0);
    REP(i, N) {
      ll lb = max(i - A[i], 0ll);
      ll ub = min(i + A[i], N - 1);
      seg.update(lb, ub + 1, 1);
    }

    bool ok = true;
    REP(i, N) {
      A[i] = seg[i];
      if (A[i] != N)
        ok = false;
    }
    if (ok)
      break;
  }

  REP(i, N) cout << A[i] << " ";
  cout << endl;

  return 0;
}