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

#include <algorithm>
#include <bitset>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>
using namespace std;
typedef long long ll;
typedef pair<ll, ll> P;
typedef pair<ll, P> P1;
typedef pair<P, P> P2;
#define pu push
#define pb push_back
#define mp make_pair
#define eps 1e-7
#define INF 1000000000
#define fi first
#define sc second
#define rep(i, x) for (ll i = 0; i < x; i++)
#define repn(i, x) for (ll i = 1; i <= x; i++)
#define SORT(x) sort(x.begin(), x.end())
#define ERASE(x) x.erase(unique(x.begin(), x.end()), x.end())
#define POSL(x, v) (lower_bound(x.begin(), x.end(), v) - x.begin())
#define POSU(x, v) (upper_bound(x.begin(), x.end(), v) - x.begin())

const int MAX = 510000;
const int MOD = 1000000007;

int main() {
  ll N, K;
  cin >> N >> K;
  vector<ll> a(N + 1);
  rep(i, N) cin >> a[i];

  vector<P> b(N);
  rep(i, N) {
    b[i].fi = max((ll)0, i - a[i]);
    b[i].sc = min(N, i + a[i] + 1);
    // cout << b[i].fi << " " << b[i].sc << endl;
  }
  SORT(b);

  vector<ll> ans(N);
  vector<ll> s(N);
  vector<ll> g(N);
  rep(i, N) s[i] = b[i].fi;
  rep(i, N) g[i] = b[i].sc;
  rep(loop, K) {
    // cout << loop << " " << ans[0] << endl;
    if (ans[0] > N - 1 and ans[N - 1] > N - 1) {
      rep(i, N) { ans[i] = N; }
      break;
    } else {
      rep(i, N) ans[i] = 0;
      ll tmp = 0;
      vector<ll> cnt(N + 4);
      // ansの更新
      rep(i, N) { cnt[g[i]]++; }
      rep(i, N) {
        tmp += POSU(s, i) - POSL(s, i);
        tmp -= cnt[i];
        ans[i] += tmp;
      }
      // sとgの更新
      rep(i, N) {
        b[i].fi = max((ll)0, i - ans[i]);
        b[i].sc = min(N, i + ans[i] + 1);
        // cout << b[i].fi << " " << b[i].sc << endl;
      }
      SORT(b);
      rep(i, N) s[i] = b[i].fi;
      rep(i, N) g[i] = b[i].sc;
    }
  }
  rep(i, N) { cout << ans[i] << " "; }
  cout << endl;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <bitset>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>
using namespace std;
typedef long long ll;
typedef pair<ll, ll> P;
typedef pair<ll, P> P1;
typedef pair<P, P> P2;
#define pu push
#define pb push_back
#define mp make_pair
#define eps 1e-7
#define INF 1000000000
#define fi first
#define sc second
#define rep(i, x) for (ll i = 0; i < x; i++)
#define repn(i, x) for (ll i = 1; i <= x; i++)
#define SORT(x) sort(x.begin(), x.end())
#define ERASE(x) x.erase(unique(x.begin(), x.end()), x.end())
#define POSL(x, v) (lower_bound(x.begin(), x.end(), v) - x.begin())
#define POSU(x, v) (upper_bound(x.begin(), x.end(), v) - x.begin())

const int MAX = 510000;
const int MOD = 1000000007;

int main() {
  ll N, K;
  cin >> N >> K;
  vector<ll> a(N + 1);
  rep(i, N) cin >> a[i];

  vector<P> b(N);
  rep(i, N) {
    b[i].fi = max((ll)0, i - a[i]);
    b[i].sc = min(N, i + a[i] + 1);
    // cout << b[i].fi << " " << b[i].sc << endl;
  }
  SORT(b);

  vector<ll> ans(N);
  vector<ll> s(N);
  vector<ll> g(N);
  rep(i, N) s[i] = b[i].fi;
  rep(i, N) g[i] = b[i].sc;
  rep(loop, K) {
    rep(i, N) ans[i] = 0;
    // cout << loop << " " << ans[0] << endl;
    if (ans[0] > N - 1 and ans[N - 1] > N - 1) {
      rep(i, N) { ans[i] = N; }
      break;
    } else {
      ll tmp = 0;
      vector<ll> cnt(N + 4);
      // ansの更新
      rep(i, N) { cnt[g[i]]++; }
      rep(i, N) {
        tmp += POSU(s, i) - POSL(s, i);
        tmp -= cnt[i];
        ans[i] += tmp;
      }
      // sとgの更新
      rep(i, N) {
        b[i].fi = max((ll)0, i - ans[i]);
        b[i].sc = min(N, i + ans[i] + 1);
        // cout << b[i].fi << " " << b[i].sc << endl;
      }
      SORT(b);
      rep(i, N) s[i] = b[i].fi;
      rep(i, N) g[i] = b[i].sc;
    }
  }
  rep(i, N) { cout << ans[i] << " "; }
  cout << endl;
}