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

#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>
using namespace std;
#define rep(i, n) for (int(i) = 0; (i) < (n); (i)++)
#define FOR(i, m, n) for (int(i) = (m); (i) < (n); (i)++)
#define All(v) (v).begin(), (v).end()
#define pb push_back
#define MP(a, b) make_pair((a), (b))
using ll = long long;
using pii = pair<int, int>;
using pll = pair<ll, ll>;
const int INF = 1 << 30;
const ll LINF = 1LL << 60;
const int MOD = 1e9 + 7;

// Range Sum Query by LazySegmentTree
struct LazySegTree {
private:
  int n;
  vector<ll> node, lazy;

public:
  LazySegTree(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];
  }

  // k番目のノードについて遅延評価
  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;
  }

  ll at(int k) { return getsum(k, k + 1); }

  // [a,b)区間加算
  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;
    }
  }

  // [a,b)区間取得
  ll getsum(int a, int b, int k = 0, int l = 0, int r = -1) {
    if (r < 0)
      r = n;
    if (b <= l || r <= a)
      return 0;

    // まず評価
    eval(k, l, r);
    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;
  }
};

int main() {
  int N, K;
  cin >> N >> K;
  vector<int> A(N);
  rep(i, N) cin >> A[i];
  rep(i, K) {
    vector<ll> imos(N + 1, 0);
    rep(j, N) {
      imos[max(0, j - A[j])]++;
      imos[min(N, j + A[j] + 1)]--;
    }
    rep(j, N) imos[j + 1] += imos[j];
    int ok = 1;
    rep(j, N) {
      if (imos[j] != N)
        ok = 0;
    }
    rep(i, N) A[i] = imos[i];
    if (ok)
      break;
  }
  rep(i, N) { cout << A[i] << (i == N - 1 ? "\\n" : " "); }
  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>
using namespace std;
#define rep(i, n) for (int(i) = 0; (i) < (n); (i)++)
#define FOR(i, m, n) for (int(i) = (m); (i) < (n); (i)++)
#define All(v) (v).begin(), (v).end()
#define pb push_back
#define MP(a, b) make_pair((a), (b))
using ll = long long;
using pii = pair<int, int>;
using pll = pair<ll, ll>;
const int INF = 1 << 30;
const ll LINF = 1LL << 60;
const int MOD = 1e9 + 7;

// Range Sum Query by LazySegmentTree
struct LazySegTree {
private:
  int n;
  vector<ll> node, lazy;

public:
  LazySegTree(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];
  }

  // k番目のノードについて遅延評価
  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;
  }

  ll at(int k) { return getsum(k, k + 1); }

  // [a,b)区間加算
  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;
    }
  }

  // [a,b)区間取得
  ll getsum(int a, int b, int k = 0, int l = 0, int r = -1) {
    if (r < 0)
      r = n;
    if (b <= l || r <= a)
      return 0;

    // まず評価
    eval(k, l, r);
    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;
  }
};

int main() {
  int N, K;
  cin >> N >> K;
  vector<ll> A(N);
  rep(i, N) cin >> A[i];
  rep(i, K) {
    vector<ll> new_A(N, 0);
    LazySegTree segTree(new_A);
    rep(j, N) { segTree.add(max(0LL, j - A[j]), min(j + A[j] + 1, (ll)N), 1); }
    int ok = 1;
    rep(j, N) {
      new_A[j] = segTree.at(j);
      if (new_A[j] != N)
        ok = 0;
    }
    A = new_A;

    if (ok)
      break;
  }
  rep(i, N) { cout << A[i] << (i == N - 1 ? "\\n" : " "); }
  return 0;
}