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

#include <algorithm>
#include <bitset>
#include <cmath>
#include <functional>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>

#define rep(i, n) for (int i = 0; i < n; i++)
using namespace std;
typedef long long ll;
using P = pair<ll, ll>;

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;
}

const ll MOD = 1000000007;
const int INF = 1 << 30;
// const ll INF = (ll)1e18 + 1;

template <class Abel> struct BIT {
  vector<Abel> dat[2];
  Abel UNITY_SUM = 0; // to be set

  /* [1, n] */
  BIT(int n) { init(n); }
  void init(int n) {
    for (int iter = 0; iter < 2; ++iter)
      dat[iter].assign(n + 1, UNITY_SUM);
  }

  /* a, b are 1-indexed, [a, b) */
  inline void sub_add(int p, int a, Abel x) {
    for (int i = a; i < (int)dat[p].size(); i += i & -i)
      dat[p][i] = dat[p][i] + x;
  }
  inline void add(int a, int b, Abel x) {
    sub_add(0, a, x * -(a - 1));
    sub_add(1, a, x);
    sub_add(0, b, x * (b - 1));
    sub_add(1, b, x * (-1));
  }

  /* a is 1-indexed, [a, b) */
  inline Abel sub_sum(int p, int a) {
    Abel res = UNITY_SUM;
    for (int i = a; i > 0; i -= i & -i)
      res = res + dat[p][i];
    return res;
  }
  inline Abel sum(int a, int b) {
    return sub_sum(0, b - 1) + sub_sum(1, b - 1) * (b - 1) - sub_sum(0, a - 1) -
           sub_sum(1, a - 1) * (a - 1);
  }
  /* get_max value */
  Abel get_max_val(int a, int b) {
    ll res = 0;
    for (int i = a; i <= b; ++i) {
      res = max(res, sum(i, i + 1));
    }
    return res;
  }

  /* debug */
  void print() {
    for (int i = 1; i < (int)dat[0].size(); ++i)
      cout << sum(i, i + 1) << ",";
    cout << endl;
  }
};

int main() {
  int N, K;
  cin >> N >> K;
  vector<int> A(N);
  rep(i, N) cin >> A[i];
  vector<int> tmp;
  BIT<ll> bit(N);
  rep(k, K) {
    /*
    for (int i = 1; i <= N; ++i) bit.add(i, i + 1, -bit.sum(i, i + 1));
    for (int i = 1; i <= N; ++i) bit.add(i, i + 1, 1);
    for (int i = 1; i <= N; ++i){
      //cout << i - A[i - 1] << " " << i + A[i - 1] + 1 << endl;
      bit.add(max(1, i - A[i - 1]), min(N + 1, i + A[i - 1] + 1), 1);
      bit.add(i, i + 1, -1);
    }*/
    for (int i = 1; i <= N; ++i)
      bit.add(i, i + 1, -bit.sum(i, i + 1));
    for (int i = 1; i <= N; ++i) {
      bit.add(max(1, i - A[i - 1]), min(N + 1, i + A[i - 1] + 1), 1);
    }
    tmp = A;
    for (int i = 1; i <= N; ++i) {
      A[i - 1] = bit.sum(i, i + 1);
    }
    if (A == tmp)
      break;
  }
  for (int i = 1; i <= N; ++i) {
    cout << bit.sum(i, i + 1) << " ";
  }
  cout << endl;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <bitset>
#include <cmath>
#include <functional>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>

#define rep(i, n) for (int i = 0; i < n; i++)
using namespace std;
typedef long long ll;
using P = pair<ll, ll>;

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;
}

const ll MOD = 1000000007;
const int INF = 1 << 30;
// const ll INF = (ll)1e18 + 1;

template <class Abel> struct BIT {
  vector<Abel> dat[2];
  Abel UNITY_SUM = 0; // to be set

  /* [1, n] */
  BIT(int n) { init(n); }
  void init(int n) {
    for (int iter = 0; iter < 2; ++iter)
      dat[iter].assign(n + 1, UNITY_SUM);
  }

  /* a, b are 1-indexed, [a, b) */
  inline void sub_add(int p, int a, Abel x) {
    for (int i = a; i < (int)dat[p].size(); i += i & -i)
      dat[p][i] = dat[p][i] + x;
  }
  inline void add(int a, int b, Abel x) {
    sub_add(0, a, x * -(a - 1));
    sub_add(1, a, x);
    sub_add(0, b, x * (b - 1));
    sub_add(1, b, x * (-1));
  }

  /* a is 1-indexed, [a, b) */
  inline Abel sub_sum(int p, int a) {
    Abel res = UNITY_SUM;
    for (int i = a; i > 0; i -= i & -i)
      res = res + dat[p][i];
    return res;
  }
  inline Abel sum(int a, int b) {
    return sub_sum(0, b - 1) + sub_sum(1, b - 1) * (b - 1) - sub_sum(0, a - 1) -
           sub_sum(1, a - 1) * (a - 1);
  }
  /* get_max value */
  Abel get_max_val(int a, int b) {
    ll res = 0;
    for (int i = a; i <= b; ++i) {
      res = max(res, sum(i, i + 1));
    }
    return res;
  }

  /* debug */
  void print() {
    for (int i = 1; i < (int)dat[0].size(); ++i)
      cout << sum(i, i + 1) << ",";
    cout << endl;
  }
};

int main() {
  int N, K;
  cin >> N >> K;
  vector<int> A(N);
  rep(i, N) cin >> A[i];
  vector<int> tmp;
  BIT<ll> bit(N);
  rep(k, K) {
    for (int i = 1; i <= N; ++i)
      bit.add(i, i + 1, -bit.sum(i, i + 1));
    for (int i = 1; i <= N; ++i)
      bit.add(i, i + 1, 1);
    for (int i = 1; i <= N; ++i) {
      // cout << i - A[i - 1] << " " << i + A[i - 1] + 1 << endl;
      bit.add(max(1, i - A[i - 1]), min(N + 1, i + A[i - 1] + 1), 1);
      bit.add(i, i + 1, -1);
    }
    tmp = A;
    for (int i = 1; i <= N; ++i) {
      A[i - 1] = bit.sum(i, i + 1);
    }
    if (A == tmp)
      break;
  }
  for (int i = 1; i <= N; ++i) {
    cout << bit.sum(i, i + 1) << " ";
  }
  cout << endl;
}