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

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>

#define pb push_back
#define F first
#define S second
#define lint long long
#define mp make_pair
#define pii pair<int, int>
#define vi vector<int>
#define sz(x) (int)x.size()
#define all(x) begin(x), end(x)
#define SET(x, val) memset(x, val, sizeof(x))
#define fastio                                                                 \\
  ios_base::sync_with_stdio(0);                                                \\
  cin.tie(0);                                                                  \\
  cout.tie(0)
using namespace std;
using namespace __gnu_pbds;
template <typename T>
using ordered_set =
    tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;

const int N = 2e5 + 5, MOD1 = 1e9 + 7, MOD2 = 998244353;
const lint INF = 1e18;

int a[N], b[N], n;
priority_queue<int, vector<int>, greater<int>> st;
priority_queue<int> st2;

bool solve() {
  while (st.size())
    st.pop();
  for (int i = 1; i <= n; i++) {
    st.push(i + a[i]);
    b[i] = st.size();
    while (st.size() && st.top() == i)
      st.pop();
  }
  /* for(int i = 1; i <= n; i++)
   {
       cout << b[i] << " ";
   }
   cout << endl;*/
  while (st2.size())
    st2.pop();
  for (int i = n; i >= 1; i--) {
    st2.push(i - a[i]);
    b[i] += st2.size() - 1;
    while (st2.size() && st2.top() == i)
      st2.pop();
  }
  int cnt = 0;
  for (int i = 1; i <= n; i++) {
    a[i] = b[i];
    if (a[i] == n)
      cnt++;
  }
  if (cnt == n)
    return true;
  return false;
}
signed main() {
  fastio;
  int k;
  cin >> n >> k;
  // n = 200000, k = 100;
  // a[1]=1;
  for (int i = 1; i <= n; i++) {
    cin >> a[i];
  }
  for (int i = 1; i <= k; i++) {
    // cout << i << endl;
    bool comp = solve();
    if (comp)
      break;
  }
  for (int i = 1; i <= n; i++) {
    cout << a[i] << " ";
  }
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>

#define pb push_back
#define F first
#define S second
#define int long long
#define lint long long
#define mp make_pair
#define pii pair<int, int>
#define vi vector<int>
#define sz(x) (int)x.size()
#define all(x) begin(x), end(x)
#define SET(x, val) memset(x, val, sizeof(x))
#define fastio                                                                 \\
  ios_base::sync_with_stdio(0);                                                \\
  cin.tie(0);                                                                  \\
  cout.tie(0)
using namespace std;
using namespace __gnu_pbds;
template <typename T>
using ordered_set =
    tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;

const int N = 2e5 + 5, MOD1 = 1e9 + 7, MOD2 = 998244353;
const lint INF = 1e18;

int a[N], b[N], n;
ordered_set<pii> st;
bool solve() {
  st.clear();
  for (int i = 1; i <= n; i++) {
    b[i] = 0;
    st.insert({i + a[i], i});
    int curr = st.size() - st.order_of_key({i, 0});
    b[i] += curr;
    while (!st.empty() && (*st.find_by_order(0)).F == i) {
      st.erase(st.find_by_order(0));
    }
  }
  /*for(int i = 1; i <= n; i++)
  {
      cout << b[i] << " ";
  }
  cout << endl;*/
  st.clear();
  for (int i = n; i >= 1; i--) {
    st.insert({-i + a[i], i});
    int curr = st.size() - st.order_of_key({-i, 0});
    b[i] += curr - 1;
    while (!st.empty() && (*st.find_by_order(0)).F == -i) {
      st.erase(st.find_by_order(0));
    }
  }
  int cnt = 0;
  for (int i = 1; i <= n; i++) {
    a[i] = b[i];
    if (a[i] == n)
      cnt++;
  }
  if (cnt == n)
    return true;
  return false;
  /*forward pass
  for(int i = 0; i < n0.size(); i++)
  {
      //beech ka
      for(int j = n0[i-1]+1; j < n0[i]; j++)
      {
          int curr = st.order_of_key({j+1, 0});
          b[j]+=curr;
          while(!st.empty() && (*st.find_by_order(0)).F==j)
          {
              st.erase(st.find_by_order(0));
          }
          if(st.empty()) break;
      }
      //iska
      b[n0[i]] = st.order_of_key({n0[i]+1, 0});
      while(!st.empty() && (*st.find_by_order(0)).F==n0[i])
      {
          st.erase(st.find_by_order(0));
      }
  }*/
}
signed main() {
  int k;
  cin >> n >> k;
  for (int i = 1; i <= n; i++) {
    cin >> a[i];
  }
  for (int i = 1; i <= k; i++) {
    bool comp = solve();
    if (comp)
      break;
  }
  for (int i = 1; i <= n; i++) {
    cout << a[i] << " ";
  }
}