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

#include <algorithm>
#include <iostream>

#ifndef DUAL_SEGMENT_TREE_HPP
#define DUAL_SEGMENT_TREE_HPP

#include <cstdint>
#include <vector>

//===
template <class Monoid> struct DualSegmentTree {
  using T = typename Monoid::value_type;
  std::vector<T> lazy;

  DualSegmentTree() = default;
  explicit DualSegmentTree(uint32_t n) : lazy(n << 1, Monoid::identity()){};

  inline int size() { return (int)(lazy.size() >> 1); };

  inline void propagate(uint32_t k) {
    if (k >= size())
      return;
    lazy[(k << 1) | 0] = Monoid::operation(lazy[(k << 1) | 0], lazy[k]);
    lazy[(k << 1) | 1] = Monoid::operation(lazy[(k << 1) | 1], lazy[k]);
    lazy[k] = Monoid::identity();
  };
  inline void push_down(uint32_t k) {
    for (uint32_t i = 31; i > 0; i--)
      propagate(k >> i);
  };

  // [l, r)
  void update(uint32_t l, uint32_t r, T op) {
    l += size();
    r += size();
    push_down(l);
    push_down(r - 1);

    while (l < r) {
      if (l & 1)
        lazy[l] = Monoid::operation(lazy[l], op), l++;
      if (r & 1)
        --r, lazy[r] = Monoid::operation(lazy[r], op);

      l >>= 1;
      r >>= 1;
    }
  };

  T operator[](uint32_t k) {
    k += size();
    push_down(k);
    return lazy[k];
  };
};
//===
#endif

using namespace std;
using llong = long long;

llong n, k;
llong a[200005];
llong imos[200005];

struct Operator {
  using value_type = llong;
  inline static llong identity() { return 0; };
  inline static llong operation(llong a, llong b) { return a + b; };
};

bool check() {
  for (int i = 0; i < n; i++) {
    if (a[i] < n)
      return false;
  }
  return true;
}

int main() {
  cin >> n >> k;
  for (int i = 0; i < n; i++) {
    cin >> a[i];
  }

  for (int j = 0; j < k; j++) {
    if (check()) {
      break;
    }

    DualSegmentTree<Operator> seg(n);

    for (int i = 0; i < n; i++) {
      seg.update(max(0ll, i - a[i]), min(n, i + a[i] + 1), 1);
    }

    for (int i = 0; i < n; i++) {
      a[i] = seg[i];
    }
  }

  for (int i = 0; i < n; i++) {
    cout << a[i] << ' ';
  }
}

Error: TLE

Fixed submission:

#pragma GCC optimize("O3")
#pragma GCC optimize("unroll-loops")

#include <algorithm>
#include <iostream>

#ifndef LAZY_SEGMENT_TREE_HPP
#define LAZY_SEGMENT_TREE_HPP

#include <cstdint>
#include <vector>
#ifndef MSB_HPP
#define MSB_HPP

#include <cstdint>

inline int msb32(uint32_t x) {
  if (x == 0)
    return -1;
#ifndef __has_builtin
  int r = 0;
  if (x >> 16)
    x >>= 16, r |= 16;
  if (x >> 8)
    x >>= 8, r |= 8;
  if (x >> 4)
    x >>= 4, r |= 4;
  if (x >> 2)
    x >>= 2, r |= 2;
  return r | (x >> 1);
#else
  return 31 - __builtin_clz(x);
#endif
};

inline int msb64(uint64_t x) {
  if (x == 0)
    return -1;

#ifndef __has_builtin
  int r = 0;
  if (x >> 32)
    x >>= 32, r |= 32;
  if (x >> 16)
    x >>= 16, r |= 16;
  if (x >> 8)
    x >>= 8, r |= 8;
  if (x >> 4)
    x >>= 4, r |= 4;
  if (x >> 2)
    x >>= 2, r |= 2;
  return r | (x >> 1);
#else
  return 63 - __builtin_clzll(x);
#endif
};

#endif

//===
template <class MonoidwithOperator> struct LazySegmentTree {
  using M = MonoidwithOperator;
  using V = typename M::value_structure;
  using T = typename V::value_type;
  using O = typename M::operator_structure;
  using E = typename O::value_type;

  // mergeMonoid V::operation
  // applyOperator M::operation
  // mergeOperator O::operation

  struct Node {
    T dat;
    E lazy;
    Node(T dat, E lazy) : dat(dat), lazy(lazy){};
  };

  std::vector<Node> tree;

  LazySegmentTree() = default;
  explicit LazySegmentTree(uint32_t n)
      : tree(n * 2 + 2, Node(V::identity(), O::identity())){};

  int size() { return tree.size() >> 1; };

  void propagation(uint32_t k) {
    const uint32_t l = (k << 1) | 0;
    const uint32_t r = (k << 1) | 1;
    tree[l].lazy = O::operation(tree[l].lazy, tree[k].lazy);
    tree[r].lazy = O::operation(tree[r].lazy, tree[k].lazy);
    tree[l].dat = M::operation(tree[l].dat, tree[k].lazy);
    tree[r].dat = M::operation(tree[r].dat, tree[k].lazy);
    tree[k].lazy = O::identity();
  };
  void push_down(uint32_t k) {
    if (k == 0)
      return;
    uint32_t w = msb32(k);
    for (int i = w; i > 0; i--)
      propagation(k >> i);
  };
  void recalc(uint32_t k) {
    while (k > 1) {
      k >>= 1;
      tree[k].dat =
          V::operation(tree[(k << 1) | 0].dat, tree[(k << 1) | 1].dat);
    }
  };

  // [l, r) += op
  void update(uint32_t l, uint32_t r, E op) {
    l += size();
    r += size();
    uint32_t tmpl = l;
    uint32_t tmpr = r;
    push_down(l);
    push_down(r - 1);

    while (l < r) {
      if (l & 1) {
        tree[l].lazy = O::operation(tree[l].lazy, op);
        tree[l].dat = M::operation(tree[l].dat, op);
        l++;
      }
      if (r & 1) {
        --r;
        tree[r].lazy = O::operation(tree[r].lazy, op);
        tree[r].dat = M::operation(tree[r].dat, op);
      }
      l >>= 1;
      r >>= 1;
    }

    push_down(tmpl);
    push_down(tmpr - 1);
    recalc(tmpl);
    recalc(tmpr - 1);
  };
  void update(uint32_t idx, T x) {
    idx += size();
    push_down(idx);
    tree[idx].dat = x;
    recalc(idx);
  };
  void set(uint32_t idx, T x) { update(idx, x); };

  // foldl[l, r)
  T fold(uint32_t l, uint32_t r) {
    l += size();
    r += size();
    push_down(l);
    push_down(r - 1);

    T lv = V::identity();
    T rv = V::identity();

    while (l < r) {
      if (l & 1)
        lv = V::operation(lv, tree[l].dat), l++;
      if (r & 1)
        --r, rv = V::operation(tree[r].dat, rv);

      l >>= 1;
      r >>= 1;
    }

    return V::operation(lv, rv);
  };

  T operator[](const uint32_t &k) {
    push_down(k + size());
    return tree[k + size()].dat;
  };
};
//===
#endif

using namespace std;
using llong = long long;

llong n, k;
llong a[200005];
llong imos[200005];

struct A {
  struct Monoid {
    using value_type = pair<llong, llong>;
    using T = value_type;
    inline static T identity() { return T(0ll, 0ll); };
    inline static T operation(T a, T b) {
      return T(a.first + b.first, a.second + b.second);
    };
  };
  struct Operator {
    using value_type = llong;
    inline static llong identity() { return 0; };
    inline static llong operation(llong a, llong b) { return a + b; };
  };
  using value_structure = Monoid;
  using operator_structure = Operator;
  using T = typename Monoid::T;
  inline static T operation(T a, llong b) {
    return T(a.first + a.second * b, a.second);
  };
};

bool check() {
  for (int i = 0; i < n; i++) {
    if (a[i] < n)
      return false;
  }
  return true;
}

int main() {
  cin >> n >> k;
  for (int i = 0; i < n; i++) {
    cin >> a[i];
  }

  for (int j = 0; j < k; j++) {
    if (check()) {
      break;
    }

    LazySegmentTree<A> seg(n);
    for (int i = 0; i < n; i++)
      seg.set(i, {0ll, 1ll});

    for (int i = 0; i < n; i++) {
      seg.update(max(0ll, i - a[i]), min(n, i + a[i] + 1), 1);
    }

    for (int i = 0; i < n; i++) {
      a[i] = seg[i].first;
    }
  }

  for (int i = 0; i < n; i++) {
    cout << a[i] << ' ';
  }
}