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

#include <bits/stdc++.h>

using namespace std;
typedef long long ll;
typedef long double ld;
typedef unsigned long long ull;
typedef pair<ll, ll> pii;
#define rep(i, x, y) for (int i = x; i < y; i++)
#define rept(i, x, y) for (int i = x; i <= y; i++)
#define per(i, x, y) for (int i = x; i >= y; i--)
#define all(x) x.begin(), x.end()
#define pb push_back
#define fi first
#define se second
#define mes(a, b) memset(a, b, sizeof a)
#define mp make_pair
#define dd(x) cout << #x << "=" << x << " "
#define de(x) cout << #x << "=" << x << "\\n"
#define debug() cout << "I love Miyamizu Mitsuha forever.\\n"
const int inf = 0x3f3f3f3f;

const int maxn = 2e5 + 5;
int a[maxn];
int cnt1[maxn], cnt2[maxn];

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  int n, k;
  cin >> n >> k;
  rept(i, 1, n) cin >> a[i];

  while (k--) {
    bool flag = 1;
    int cnt = 0;
    mes(cnt1, 0);
    mes(cnt2, 0);
    rept(i, 1, n) {
      cnt1[max(1, i - a[i])]++;
      cnt2[min(n + 1, i + a[i] + 1)]++;
    }
    rept(i, 1, n) {
      cnt += cnt1[i] - cnt2[i];
      if (cnt != n)
        flag = 0;
      a[i] = cnt;
    }
    if (flag)
      break;
  }
  rept(i, 1, n) cout << a[i] << " ";
  cout << "\\n";
  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;
typedef long long ll;
typedef long double ld;
typedef unsigned long long ull;
typedef pair<ll, ll> pii;
#define rep(i, x, y) for (int i = x; i < y; i++)
#define rept(i, x, y) for (int i = x; i <= y; i++)
#define per(i, x, y) for (int i = x; i >= y; i--)
#define all(x) x.begin(), x.end()
#define pb push_back
#define fi first
#define se second
#define mes(a, b) memset(a, b, sizeof a)
#define mp make_pair
#define dd(x) cout << #x << "=" << x << " "
#define de(x) cout << #x << "=" << x << "\\n"
#define debug() cout << "I love Miyamizu Mitsuha forever.\\n"
const int inf = 0x3f3f3f3f;

const int maxn = 2e5 + 5;
int a[maxn], b[maxn];

class element {
public:
  int l, r;
  int plus, val;
};

class Tree {
public:
  element tree[maxn << 2];
  void build(int id, int l, int r) {
    tree[id].l = l;
    tree[id].r = r;
    tree[id].plus = 0;
    if (l == r) {
      tree[id].val = 0;
      return;
    }
    int mid = (l + r) / 2;
    build(id * 2, l, mid);
    build(id * 2 + 1, mid + 1, r);
    tree[id].val = tree[id * 2].val + tree[id * 2 + 1].val;
  }
  void add(int id, int l, int r, int num) {
    if (tree[id].l >= l && tree[id].r <= r) {
      tree[id].plus += num;
      tree[id].val += num * (tree[id].r - tree[id].l + 1);
      return;
    }
    if (tree[id].l > r || tree[id].r < l)
      return;
    push_down(id);
    if (tree[id * 2].r >= l)
      add(id * 2, l, r, num);
    if (tree[id * 2 + 1].l <= r)
      add(id * 2 + 1, l, r, num);
    tree[id].val = tree[id * 2].val + tree[id * 2 + 1].val;
  }
  void push_down(int id) {
    if (tree[id].l != tree[id].r)
      rept(j, id * 2, id * 2 + 1) {
        tree[j].val += tree[id].plus * (tree[j].r - tree[j].l + 1);
        tree[j].plus += tree[id].plus;
      }
    tree[id].plus = 0;
  }
  int query(int id, int l, int r) {
    if (tree[id].l >= l && tree[id].r <= r)
      return tree[id].val;
    if (tree[id].l > r || tree[id].r < l)
      return 0;
    ll s = 0;
    push_down(id);
    if (tree[id * 2].r >= l)
      s += query(id * 2, l, r);
    if (tree[id * 2 + 1].l <= r)
      s += query(id * 2 + 1, l, r);
    return s;
  }
} t;

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  int n, k;
  cin >> n >> k;
  if (k >= 42) {
    rept(i, 1, n) cout << n << " ";
    cout << "\\n";
    return 0;
  }
  t.build(1, 1, n);
  rept(i, 1, n) cin >> a[i];
  //	rept(i,1,n) a[i]=0;
  rept(i, 1, n) t.add(1, i - a[i], i + a[i], 1);

  while (k--) {
    bool flag = 1;
    rept(i, 1, n) {
      b[i] = t.query(1, i, i);
      if (b[i] != n)
        flag = 0;
    }
    rept(i, 1, n) {
      t.add(1, i - a[i], i + a[i], -1);
      t.add(1, i - b[i], i + b[i], 1);
    }
    swap(a, b);
    if (flag)
      break;
  }
  rept(i, 1, n) cout << a[i] << " ";
  cout << "\\n";
  return 0;
}