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

#include <bits/stdc++.h>
using namespace std;
#ifdef LOCAL
#include "debug.h"
#else
#define dump(...)
#endif
#define endl "\\n"
#define ll long long
#define FOR(i, a, b) for (int i = (int)(a); i < (int)(b); i++)
#define REP(i, x) for (int i = 0; i < (int)(x); i++)
#define REPS(i, x) for (int i = 1; i <= (int)(x); i++)
#define RREP(i, x) for (int i = ((int)(x)-1); i >= 0; i--)
#define RREPS(i, x) for (int i = ((int)(x)); i > 0; i--)
#define INF (1 << 30 - 1)
#define LLINF (1LL << 60)
#define vi vector<int>
#define vvi vector<vector<int>>
#define pi pair<int, int>
#define ALL(a) (a).begin(), (a).end()
#define UNIQUE(v)                                                              \\
  v.erase(unique(v.begin(), v.end()), v.end()) // sortしてからつかうこと
constexpr ll MOD = 1e9 + 7;
int dx[8] = {1, 0, -1, 0, 1, 1, -1, -1};
int dy[8] = {0, 1, 0, -1, 1, -1, 1, -1};
ll A, B, C, D, E, F, G, H, N, M, L, K, P, Q, R, W, X, Y, Z;
ll V, T;
string S;
ll ans = 0;

template <typename T> istream &operator>>(istream &is, vector<T> &vec) {
  for (T &x : vec)
    is >> x;
  return is;
}

template <typename T> class BIT {
public:
  int n;            // 要素数
  vector<T> bit[2]; // データの格納先
  BIT(int n_) { init(n_); }
  void init(int n_) {
    n = n_ + 1;
    for (int p = 0; p < 2; p++)
      bit[p].assign(n, 0);
  }

  void add_sub(int p, int i, T x) {
    for (int idx = i; idx < n; idx += (idx & -idx)) {
      bit[p][idx] += x;
    }
  }
  void add(int l, int r, T x) { // [l,r) に加算
    add_sub(0, l, -x * (l - 1));
    add_sub(0, r, x * (r - 1));
    add_sub(1, l, x);
    add_sub(1, r, -x);
  }

  T sum_sub(int p, int i) {
    T s(0);
    for (int idx = i; idx > 0; idx -= (idx & -idx)) {
      s += bit[p][idx];
    }
    return s;
  }
  T sum(int i) { return sum_sub(0, i) + sum_sub(1, i) * i; }
};

signed main() {
  cin >> N >> K;
  vector<ll> a(N);
  cin >> a;
  int cnt = 0;
  int k = 0;
  while (cnt != N && k < K) {
    cnt = 0;
    BIT<int> bt(N);
    k++;
    REP(i, N) {
      int j = i + 1;
      int l = max(1LL, j - a[i]);
      int r = min(N + 1, j + a[i] + 1);
      bt.add(l, r, 1);
    }
    REP(i, N) { a[i] = bt.sum(i + 1) - bt.sum(i); }
    REP(i, N) {
      if (a[i] == N) {
        cnt++;
      }
    }
  }
  REP(i, N) { cout << a[i] << " "; }
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#ifdef LOCAL
#include "debug.h"
#else
#define dump(...)
#endif
#define endl "\\n"
#define ll long long
#define FOR(i, a, b) for (int i = (int)(a); i < (int)(b); i++)
#define REP(i, x) for (int i = 0; i < (int)(x); i++)
#define REPS(i, x) for (int i = 1; i <= (int)(x); i++)
#define RREP(i, x) for (int i = ((int)(x)-1); i >= 0; i--)
#define RREPS(i, x) for (int i = ((int)(x)); i > 0; i--)
#define INF (1 << 30 - 1)
#define LLINF (1LL << 60)
#define vi vector<int>
#define vvi vector<vector<int>>
#define pi pair<int, int>
#define ALL(a) (a).begin(), (a).end()
#define BIT(n) (1LL << (n))
#define UNIQUE(v)                                                              \\
  v.erase(unique(v.begin(), v.end()), v.end()) // sortしてからつかうこと
constexpr ll MOD = 1e9 + 7;
int dx[8] = {1, 0, -1, 0, 1, 1, -1, -1};
int dy[8] = {0, 1, 0, -1, 1, -1, 1, -1};
ll A, B, C, D, E, F, G, H, N, M, L, K, P, Q, R, W, X, Y, Z;
ll V, T;
string S;
ll ans = 0;

template <typename T> istream &operator>>(istream &is, vector<T> &vec) {
  for (T &x : vec)
    is >> x;
  return is;
}

class LazySegmentTree {
private:
  int n;
  vector<ll> node, lazy;

public:
  LazySegmentTree(vector<ll> v) {
    int sz = (int)v.size();
    n = 1;
    while (n < sz)
      n *= 2;
    node.resize(2 * n - 1);
    lazy.resize(2 * n - 1, 0);

    for (int i = 0; i < sz; i++)
      node[i + n - 1] = v[i];
    for (int i = n - 2; i >= 0; i--)
      node[i] = node[i * 2 + 1] + node[i * 2 + 2];
  }

  void eval(int k, int l, int r) {
    if (lazy[k] != 0) {
      node[k] += lazy[k];
      if (r - l > 1) {
        lazy[2 * k + 1] += lazy[k] / 2;
        lazy[2 * k + 2] += lazy[k] / 2;
      }
      lazy[k] = 0;
    }
  }

  void add(int a, int b, ll x, int k = 0, int l = 0, int r = -1) {
    if (r < 0)
      r = n;
    eval(k, l, r);
    if (b <= l || r <= a)
      return;
    if (a <= l && r <= b) {
      lazy[k] += (r - l) * x;
      eval(k, l, r);
    } else {
      add(a, b, x, 2 * k + 1, l, (l + r) / 2);
      add(a, b, x, 2 * k + 2, (l + r) / 2, r);
      node[k] = node[2 * k + 1] + node[2 * k + 2];
    }
  }

  ll getsum(int a, int b, int k = 0, int l = 0, int r = -1) {
    if (r < 0)
      r = n;
    eval(k, l, r);
    if (b <= l || r <= a)
      return 0;
    if (a <= l && r <= b)
      return node[k];
    ll vl = getsum(a, b, 2 * k + 1, l, (l + r) / 2);
    ll vr = getsum(a, b, 2 * k + 2, (l + r) / 2, r);
    return vl + vr;
  }
};

signed main() {
  cin >> N >> K;
  vector<ll> a(N);
  cin >> a;
  int cnt = 0;
  int k = 0;
  while (cnt != N && k < K) {
    cnt = 0;
    LazySegmentTree st(vector<ll>(N, 0));
    k++;
    REP(i, N) {
      int l = max(0LL, i - a[i]);
      int r = min(N, i + a[i] + 1);
      st.add(l, r, 1);
    }
    REP(i, N) { a[i] = st.getsum(i, i + 1); }
    REP(i, N) {
      if (a[i] == N) {
        cnt++;
      }
    }
  }
  REP(i, N) { cout << a[i] << " "; }
}