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

#include <bits/stdc++.h>
using namespace std;
#define int long long
#define rep(i, n) for (int i = 0; i < n; i++)
#define REP(i, n) for (int i = 1; i < n; i++)
#define rev(i, n) for (int i = n - 1; i >= 0; i--)
#define all(v) v.begin(), v.end()
#define P pair<int, int>
#define len(s) (int)s.size()
#define pb push_back

template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
constexpr int mod = 1e9 + 7;
constexpr int inf = 3e18;

int N, K, A[200005];
int n[200005];
signed main() {
  cin >> N >> K;
  rep(i, N) { cin >> A[i]; }
  rep(i, K) {
    memset(n, 0, sizeof(n));
    rep(j, N) {
      int l = j - A[j], r = j + A[j] + 1;
      chmax(l, 0ll);
      chmin(r, N);
      n[l]++;
      n[r]--;
    }
    bool flag = true;
    rep(j, N) {
      n[j + 1] += n[j];
      A[j] = n[j];
      if (n[j] != N)
        flag = false;
    }
    if (flag) {
      break;
    }
  }
  rep(i, N) cout << A[i] << " ";
  cout << endl;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#define int long long
#define rep(i, n) for (int i = 0; i < n; i++)
#define REP(i, n) for (int i = 1; i < n; i++)
#define rev(i, n) for (int i = n - 1; i >= 0; i--)
#define all(v) v.begin(), v.end()
#define P pair<int, int>
#define len(s) (int)s.size()
#define pb push_back

template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
constexpr int mod = 1e9 + 7;
constexpr int inf = 3e18;

template <typename Monoid, typename OperatorMonoid> struct Segtree {
  using F = function<Monoid(Monoid, Monoid)>;
  using G = function<Monoid(Monoid, OperatorMonoid, int)>;
  using H = function<OperatorMonoid(OperatorMonoid, OperatorMonoid)>;
  int size = 1;
  vector<Monoid> dat;
  vector<OperatorMonoid> lazy;
  const F f;
  const G g;
  const H h;
  Monoid M, M1;
  OperatorMonoid OM;
  void set(int a, Monoid x) { dat[a + size - 1] = x; }
  void init() {
    for (int i = size - 2; i >= 0; i--) {
      dat[i] = f(dat[i * 2 + 1], dat[i * 2 + 2]);
    }
  }
  void eval(int k, int l, int r) {
    if (lazy[k] != OM) {
      dat[k] = g(dat[k], lazy[k], (r - l));
      if (r - l > 1) {
        lazy[2 * k + 1] = h(lazy[2 * k + 1], lazy[k]);
        lazy[2 * k + 2] = h(lazy[2 * k + 2], lazy[k]);
      }
      lazy[k] = OM;
    }
  }
  void update(int a, int b, OperatorMonoid M, int k = 0, int l = 0,
              int r = -1) {
    if (r == -1)
      r = size;
    eval(k, l, r);
    if (r <= a || b <= l)
      return;
    if (a <= l && r <= b) {
      lazy[k] = h(lazy[k], M);
      eval(k, l, r);
      return;
    }
    update(a, b, M, k * 2 + 1, l, (l + r) / 2);
    update(a, b, M, k * 2 + 2, (l + r) / 2, r);
    dat[k] = f(dat[k * 2 + 1], dat[k * 2 + 2]);
  }
  Monoid query(int a, int b, int k = 0, int l = 0, int r = -1) {
    if (r == -1)
      r = size;
    eval(k, l, r);
    if (r <= a || b <= l)
      return M1;
    if (a <= l && r <= b)
      return dat[k];
    Monoid lv = query(a, b, k * 2 + 1, l, (l + r) / 2);
    Monoid rv = query(a, b, k * 2 + 2, (l + r) / 2, r);
    return f(lv, rv);
  }
  Segtree(int x, F f, G g, H h, Monoid M, OperatorMonoid OM, Monoid M1)
      : f(f), g(g), h(h), M(M), OM(OM), M1(M1) {
    while (size < x)
      size *= 2;
    dat.resize(size * 2 - 1, M);
    lazy.resize(size * 2 - 1, OM);
  }
};
int N, K, A[200005];
int n[200005];
signed main() {
  cin >> N >> K;
  bool flag = false;
  rep(i, N) { cin >> A[i]; }
  auto f = [](int a, int b) { return a + b; };
  auto g = [](int a, int b, int sz) { return a + b * sz; };
  auto h = [](int a, int b) { return a + b; };
  rep(i, K) {
    Segtree<int, int> segtree(N, f, g, h, 0, 0, 0);
    rep(j, N) {
      int l = j - A[j], r = j + A[j] + 1;
      chmax(l, 0ll);
      chmin(r, N);
      segtree.update(l, r, 1);
    }
    bool flag = false;
    rep(j, N) {
      int k = segtree.query(j, j + 1);
      if (k != N)
        flag = true;
      A[j] = k;
    }
    if (!flag)
      break;
  }
  rep(i, N) cout << A[i] << " ";
  cout << endl;
}