1701D - Permutation Restoration - CodeForces Solution

#include<bits/stdc++.h>
#include<math.h>
#include<ext/pb_ds/assoc_container.hpp>
#include<ext/pb_ds/tree_policy.hpp>
#define ll int
#define fr(i, a, n) for(ll i=a; i<n; i++)
#define frn(i, a, n) for(ll i=a; i>n; i--)
#define pll pair<ll, ll>
#define vll vector<ll>
#define append push_back
#define add insert
#define umap unordered_map
#define pqu priority_queue
using namespace std;
using namespace __gnu_pbds;

// typedef tree<ll, null_type, less_equal<ll>, rb_tree_tag, tree_order_statistics_node_update> pbds;
//order_of_key (k) : Number of items strictly smaller than k .
//find_by_order(k) : pointer to K-th element in a set (counting from zero).
// custom set comparator set<vector<int>, decltype(cmp)*> s1(cmp);

const ll maxi = 500005, mod = 1000000007;
ll testcc, n, k;
ll a[maxi], res[maxi];
set<pll> sets[maxi];

class Compare {
public:
	bool operator()(ll a, ll b)
	{
		return (*sets[a].begin()).first > (*sets[b].begin()).first;
	}
};

priority_queue<ll, vector<ll>, Compare> pq1, empt;






int main() {
	ios_base::sync_with_stdio(false);
	cin.tie(0);
	cout.tie(0);
#ifndef ONLINE_JUDGE
	// for getting input from input.txt
	freopen("input.txt", "r", stdin);
	// for writing output to output.txt
	freopen("output.txt", "w", stdout);
#endif

	cin >> testcc;
	fr(i_, 0, testcc) {
		// s1.clear();
		pq1 = empt;
		cin >> n;
		fr(i, 1, n + 1) {
			sets[i].clear();
			cin >> a[i];
		}

		fr(i, 1, n + 1) {
			ll p, q;
			if (a[i] == 0) {
				p = i + 1, q = n;
			}
			else {
				p = (i / (a[i] + 1)) + 1, q = i / a[i];
			}
			sets[p].add({q, i});
		}


		ll ind1 = 1, ind2 = 1;
		while (true) {

			if (ind2 > n) {
				break;
			}

			while (ind1 <= n and ((pq1.size() == 0) or ind1 <= ind2)) {
				if (sets[ind1].size() != 0) {
					pq1.push(ind1);
				}
				ind1 += 1;
			}
			while (ind2 < ind1 and pq1.size() != 0) {
				ll t = pq1.top();
				pq1.pop();

				res[(*sets[t].begin()).second] = ind2;


				sets[t].erase(sets[t].begin());
				if (sets[t].size() != 0) {
					pq1.push(t);
				}

				ind2 += 1;
				break;
			}
		}

		fr(i, 1, n + 1) {
			cout << res[i] << " ";
		}
		cout << endl;




	}

	return 0;
}