Problem p04052 (C++)

Buggy submission (status: Runtime Error):
/*input
4 2
4 2 3 1

*/

// spoiled

#include <algorithm>
#include <bitset>
#include <cassert>
#include <climits>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <iostream>
#include <numeric>
#include <queue>
#include <string>
#include <vector>

using namespace std;

#define FOR(i, a, b) for (int i = a; i <= b; ++i)
#define FORA(i, a) for (auto &i : a)
#define FORB(i, a, b) for (int i = a; i >= b; --i)
#define SZ(a) ((int)a.size())
#define ALL(a) begin(a), end(a)

typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
typedef vector<int> vi;
typedef vector<ll> vl;
#define fi first
#define se second

// start of code

const int MAXN = 500010, LOG = 20;

int fenwick[MAXN];
void upd(int i, int x) {
  for (; i < MAXN; i += i & -i)
    fenwick[i] += x;
}
int get(int i) {
  int ans = 0;
  for (; i; i -= i & -i)
    ans += fenwick[i];
  return ans;
}
int sch(int x) {
  int i = 0;
  for (int lg = LOG - 1; lg >= 0; --lg) {
    if (fenwick[i + (1 << lg)] < x)
      i += 1 << lg, x -= fenwick[i];
  }
  return i + 1;
}

int N, K;
int outDeg[MAXN], P[MAXN], I[MAXN];
priority_queue<int> pq;
vector<int> rG[MAXN];

signed main() {
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  cout.tie(0);

  cin >> N >> K;

  for (int i = 0; i < N; ++i) {
    cin >> P[i];
    I[P[i]] = i;

    upd(P[i], 1);
    if (i >= K)
      upd(P[i - K], -1);

    int rnk = get(P[i]);

    if (rnk > 1) {
      int j = I[sch(rnk - 1)];
      rG[i].push_back(j);
      ++outDeg[j];
      // cout << j << " -> " << i << endl;
    }

    if (rnk < min(i + 1, K)) {
      int j = I[sch(rnk + 1)];
      rG[j].push_back(i);
      ++outDeg[i];
      // cout << i << " -> " << j << endl;
    }
  }

  for (int i = 0; i < N; ++i) {
    if (outDeg[i] == 0) {
      pq.emplace(i);
    }
  }

  for (int i = N; i >= 1; --i) {
    int u = pq.top();
    pq.pop();
    // cout << i << ": " << u << endl;
    P[u] = i;
    for (int v : rG[u]) {
      --outDeg[v];
      if (outDeg[v] == 0)
        pq.emplace(v);
    }
  }

  for (int i = 0; i < N; ++i) {
    cout << P[i] << '\n';
  }

  return 0;
}


Fixed submission:
/*input
4 2
4 2 3 1

*/

// spoiled

#include <algorithm>
#include <bitset>
#include <cassert>
#include <climits>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <iostream>
#include <numeric>
#include <queue>
#include <string>
#include <vector>

using namespace std;

#define FOR(i, a, b) for (int i = a; i <= b; ++i)
#define FORA(i, a) for (auto &i : a)
#define FORB(i, a, b) for (int i = a; i >= b; --i)
#define SZ(a) ((int)a.size())
#define ALL(a) begin(a), end(a)

typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
typedef vector<int> vi;
typedef vector<ll> vl;
#define fi first
#define se second

// start of code

const int MAXN = 500010, LOG = 20;

int fenwick[MAXN];
void upd(int i, int x) {
  for (; i < MAXN; i += i & -i)
    fenwick[i] += x;
}
int get(int i) {
  int ans = 0;
  for (; i; i -= i & -i)
    ans += fenwick[i];
  return ans;
}
int sch(int x) {
  int i = 0;
  for (int lg = LOG - 1; lg >= 0; --lg) {
    if (i + (1 << lg) < MAXN && fenwick[i + (1 << lg)] < x)
      i += 1 << lg, x -= fenwick[i];
  }
  return i + 1;
}

int N, K;
int outDeg[MAXN], P[MAXN], I[MAXN];
priority_queue<int> pq;
vector<int> rG[MAXN];

signed main() {
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  cout.tie(0);

  cin >> N >> K;

  for (int i = 0; i < N; ++i) {
    cin >> P[i];
    I[P[i]] = i;

    upd(P[i], 1);
    if (i >= K)
      upd(P[i - K], -1);

    int rnk = get(P[i]);

    if (rnk > 1) {
      int j = I[sch(rnk - 1)];
      rG[i].push_back(j);
      ++outDeg[j];
      // cout << j << " -> " << i << endl;
    }

    if (rnk < min(i + 1, K)) {
      int j = I[sch(rnk + 1)];
      rG[j].push_back(i);
      ++outDeg[i];
      // cout << i << " -> " << j << endl;
    }
  }

  for (int i = 0; i < N; ++i) {
    if (outDeg[i] == 0) {
      pq.emplace(i);
    }
  }

  for (int i = N; i >= 1; --i) {
    int u = pq.top();
    pq.pop();
    // cout << i << ": " << u << endl;
    P[u] = i;
    for (int v : rG[u]) {
      --outDeg[v];
      if (outDeg[v] == 0)
        pq.emplace(v);
    }
  }

  for (int i = 0; i < N; ++i) {
    cout << P[i] << '\n';
  }

  return 0;
}