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

#include <iostream>
#include <map>
#include <string>
#include <vector>
#define rep(i, start, end) for (int i = (int)start; i < (int)end; ++i)
#define rrep(i, start, end) for (int i = (int)start - 1; i >= (int)end; --i)
#define all(x) (x).begin(), (x).end()
using namespace std;
using ll = long long;
template <typename T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return 0;
}
template <typename T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return 0;
}

vector<int> change(int N, const vector<int> &A) {
  vector<int> res(N, 0);
  rep(i, 0, N) {
    ++res[max(0, i - A[i])];
    if (i + A[i] + 1 < N) {
      --res[i + A[i] + 1];
    }
  }
  rep(i, 1, N) { res[i] += res[i - 1]; }
  return res;
}

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  int N, K;
  cin >> N >> K;
  vector<int> A(N);
  for (auto &a : A) {
    cin >> a;
  }
  rep(i, 0, K) {
    vector<int> B = change(N, A);
    if (A == B) {
      break;
    }
    swap(A, B);
  }
  rep(i, 0, N) { cout << A[i] << " \\n"[i + 1 == N]; }
  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <functional>
#include <iostream>
#include <string>
#include <vector>
#define rep(i, start, end) for (int i = (int)start; i < (int)end; ++i)
#define rrep(i, start, end) for (int i = (int)start - 1; i >= (int)end; --i)
#define all(x) (x).begin(), (x).end()
using namespace std;
using ll = long long;
template <typename T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return 0;
}
template <typename T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return 0;
}

template <typename T> class SegTree {
private:
  size_t node_size_;
  vector<T> node_;
  T identity_;                  // ex) INF, 0, etc...
  function<T(T, T)> operation_; // ex) min, max, sum, etc...
  function<T(T, T)> update_;    // ex) add, change, etc..
public:
  SegTree(size_t node_size, T identity, function<T(T, T)> operation,
          function<T(T, T)> update)
      : identity_(identity), operation_(operation), update_(update) {
    node_size_ = 1;
    while (node_size_ < node_size)
      node_size_ *= 2;
    node_ = vector<T>(2 * node_size_ - 1, identity);
  }
  void change(int idx, T new_value) {
    int node_idx = idx + node_size_ - 1;
    node_[node_idx] = update_(node_[node_idx], new_value);
    while (node_idx > 0) {
      node_idx = (node_idx - 1) / 2;
      node_[node_idx] =
          operation_(node_[node_idx * 2 + 1], node_[node_idx * 2 + 2]);
    }
  }
  T _query(int left, int right, int node_idx, int node_left, int node_right) {
    if (node_right <= left || right <= node_left) {
      return identity_;
    }
    if (left <= node_left && node_right <= right) {
      return node_[node_idx];
    } else {
      T left_value = _query(left, right, node_idx * 2 + 1, node_left,
                            (node_left + node_right) / 2);
      T right_value = _query(left, right, node_idx * 2 + 2,
                             (node_left + node_right) / 2, node_right);
      return operation_(left_value, right_value);
    }
  }
  T query(int left, int right) { return _query(left, right, 0, 0, node_size_); }
  T operator[](int idx) { return node_[idx + node_size_ - 1]; }
  int _search(int kth, int node_idx, int node_left, int node_right) {
    // kth >= 1
    if (node_right - node_left == 1)
      return node_idx - node_size_ + 1;
    if (node_[node_idx * 2 + 1] < kth)
      return _search(kth - node_[node_idx * 2 + 1], node_idx * 2 + 2,
                     (node_left + node_right) / 2, node_right);
    else
      return _search(kth, node_idx * 2 + 1, node_left,
                     (node_left + node_right) / 2);
  }
  int search(int kth) { return _search(kth, 0, 0, node_size_); }
  void clear() { fill(all(node_), identity_); }
};

int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);
  int N, K;
  cin >> N >> K;
  vector<int> A(N);
  for (auto &a : A) {
    cin >> a;
  }
  SegTree<int> seg1(
      N, 0, [](int a, int b) { return a + b; },
      [](int a, int b) { return a + b; }); // i - A[i]
  SegTree<int> seg2(
      N, 0, [](int a, int b) { return a + b; },
      [](int a, int b) { return a + b; }); // i + A[i]
  rep(_, 0, K) {
    seg1.clear();
    seg2.clear();
    rep(i, 0, N) { seg1.change(max(0, i - A[i]), 1); }
    vector<int> B(N);
    rep(i, 0, N) {
      B[i] = seg1.query(0, i + 1) + seg2.query(i, N);
      seg1.change(max(0, i - A[i]), -1);
      seg2.change(min(N - 1, i + A[i]), 1);
    }
    swap(A, B);
    bool is_ok = true;
    rep(i, 0, N) {
      if (A[i] != N) {
        is_ok = false;
        break;
      }
    }
    if (is_ok) {
      break;
    }
  }
  rep(i, 0, N) { cout << A[i] << " \\n"[i + 1 == N]; }
  return 0;
}