Thieves and Prisons

input/code.cpp: In function 'int main()':
input/code.cpp:75:14: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
    if(bal[i] > Q.size() or in_prison[ev[i].nd]) nope();
       ~~~~~~~^~~~~~~~~~
input/code.cpp:76:17: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
    while(bal[i] <= Q.size()) {
          ~~~~~~~^~~~~~~~~~~

#include<bits/stdc++.h>
using namespace std;
#ifndef d
#define d(...)
#endif
#define st first
#define nd second
#define pb push_back
#define siz(c) (int)(c).size()
#define all(c) (c).begin(), (c).end()
typedef long long LL;
typedef long double LD;
constexpr int INF=1e9+7;
constexpr LL INFL=1e18;
template<class L, class R> ostream &operator<<(ostream &os, pair<L,R> P) {
  return os << "(" << P.st << "," << P.nd << ")";
}
/*
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace __gnu_pbds;
template <typename T>
using ordered_set =
    tree<T, null_type, less<T>, rb_tree_tag,
         tree_order_statistics_node_update>;
*/
 
constexpr int maxn = 100005;
 
int n, k, m, bal[maxn];
vector<pair<bool, int>> ev;
vector<int> pos[maxn];
int pt[maxn];
 
vector<int> escaped[maxn];
bool in_prison[maxn];
 
set<int> prison[maxn];
int where[maxn];
 
int ans[maxn];
 
int main()
{
	ios_base::sync_with_stdio(0);
	cin.tie(0);
	
	cin >> n >> k >> m;
	for(int i=0; i<m; i++) {
		char c; int x;
		cin >> c >> x;
		ev.emplace_back(c == 'O', x);
		pos[x].push_back(i);
	}
	
	int cur = 0;
	for(int i=m-1; i>=0; i--) {
		cur += (ev[i].st ? 1 : -1);
		cur = max(cur, 0);
		bal[i] = cur;
	}
	
	for(int i=1; i<=n; i++)
		pos[i].push_back(m);
	
	priority_queue<pair<int, int>, vector<pair<int, int>>, greater<pair<int, int>>> Q;
	
	auto nope = []() {
		cout << "IMPOSSIBLE\n";
		exit(0);
	};
	
	for(int i=0; i<m; i++) {
		if(ev[i].st) {
			if(bal[i] > Q.size() or in_prison[ev[i].nd]) nope();
			while(bal[i] <= Q.size()) {
				assert(not Q.empty());
				auto p = Q.top(); Q.pop();
				if(p.st <= i) nope();
				escaped[i].push_back(p.nd);
				in_prison[p.nd] = false;
			}
			pt[ev[i].nd]++;
		} else {
			if(in_prison[ev[i].nd]) nope();
			in_prison[ev[i].nd] = true;
			Q.emplace(pos[ev[i].nd][++pt[ev[i].nd]], ev[i].nd);
		}
	}
	
	while(not Q.empty()) {
		auto p = Q.top(); Q.pop();
		escaped[m].push_back(p.nd);
	}
	
	set<int> av;
	for(int i=1; i<=k; i++) av.emplace(i);
	for(int i=m; i>=0; i--) {
		if(not escaped[i].empty()) {
			if(av.empty()) nope();
			auto j = *av.begin();
			av.erase(av.begin());
			for(auto x:escaped[i]) 
				prison[where[x] = j].insert(x);
			ans[i] = j;
		} else if(i < m) {
			int t = ev[i].nd;
			prison[where[t]].erase(t);
			if(prison[where[t]].empty())
				av.insert(where[t]);
			ans[i] = where[t];
		}
	}
	
	for(int i=0; i<m; i++)
		cout << ans[i] << " ";
	cout << "\n";
}

Test 1

Group: 2, 4, 5

Verdict: ACCEPTED

Test 2

Group: 2, 4, 5

Verdict: ACCEPTED

Test 3

Group: 2, 4, 5

Verdict: ACCEPTED

Test 4

Group: 2, 4, 5

Verdict: ACCEPTED

Test 5

Group: 2, 4, 5

Verdict: ACCEPTED

Test 6

Group: 4, 5

Verdict: ACCEPTED

Test 7

Group: 4, 5

Verdict: ACCEPTED

Test 8

Group: 4, 5

Verdict: ACCEPTED

Test 9

Group: 1, 3, 4, 5

Verdict: ACCEPTED

Test 10

Group: 1, 3, 4, 5

Verdict: ACCEPTED

Test 11

Group: 1, 3, 4, 5

Verdict: ACCEPTED

Test 12

Group: 1, 3, 4, 5

Verdict: ACCEPTED

Test 13

Group: 1, 3, 4, 5

Verdict: ACCEPTED

Test 14

Group: 1, 3, 4, 5

Verdict: ACCEPTED

Test 15

Group: 1, 3, 4, 5

Verdict: ACCEPTED

Test 16

Group: 1, 3, 4, 5

Verdict: ACCEPTED

Test 17

Group: 1, 2, 3, 4, 5

Verdict: ACCEPTED

Test 18

Group: 1, 3, 4, 5

Verdict: ACCEPTED

Test 19

Group: 2, 5

Verdict: ACCEPTED

Test 20

Group: 2, 5

Verdict: ACCEPTED

Test 21

Group: 2, 5

Verdict: ACCEPTED

Test 22

Group: 5

Verdict: ACCEPTED

Test 23

Group: 5

Verdict: ACCEPTED

Test 24

Group: 3, 4, 5

Verdict: ACCEPTED

Test 25

Group: 3, 4, 5

Verdict: ACCEPTED

Test 26

Group: 3, 4, 5

Verdict: ACCEPTED

Test 27

Group: 3, 4, 5

Verdict: ACCEPTED

Test 28

Group: 4, 5

Verdict: ACCEPTED

Test 29

Group: 4, 5

Verdict: ACCEPTED

Test 30

Group: 4, 5

Verdict: ACCEPTED

Test 31

Group: 4, 5

Verdict: ACCEPTED

Test 32

Group: 2, 4, 5

Verdict: ACCEPTED

Test 33

Group: 2, 4, 5

Verdict: ACCEPTED

Test 34

Group: 2, 4, 5

Verdict: ACCEPTED

Test 35

Group: 2, 4, 5

Verdict: ACCEPTED

Test 36

Group: 3, 5

Verdict: ACCEPTED

Test 37

Group: 3, 5

Verdict: ACCEPTED

Test 38

Group: 3, 5

Verdict: ACCEPTED

Test 39

Group: 3, 5

Verdict: ACCEPTED

Test 40

Group: 5

Verdict: ACCEPTED

Test 41

Group: 5

Verdict: ACCEPTED

Test 42

Group: 5

Verdict: ACCEPTED

Test 43

Group: 5

Verdict: ACCEPTED

Test 44

Group: 2, 5

Verdict: ACCEPTED

Test 45

Group: 2, 5

Verdict: ACCEPTED

Test 46

Group: 2, 5

Verdict: ACCEPTED

Test 47

Group: 2, 5

Verdict: ACCEPTED