Thieves and Prisons

#include <vector>
#include <iostream>
#include <algorithm>

using namespace std;
vector<int> openings;  //sorted list of openings;
vector<int> val;
vector<int> lC;
vector<int> rC;
vector<int> parent; 

int root;



int query(int value){
	int best_ind = -1;
	int cur_best = 100000;
	int cur_node;
	int next = root;

	while(next != -1){
		cur_node = next;
		if(val[cur_node] >= value and val[cur_node] < cur_best){
			cur_best = val[cur_node];
			best_ind = cur_node;
		}
		if(val[cur_node] >= cur_best){
			next = lC[cur_node];
		}else{
			next = rC[cur_node]; 
		}
	}
	return best_ind;
}


int find_min(int node){
	int cur_node = node;
	while(lC[cur_node] != -1){
		cur_node = lC[cur_node];
	}
	return cur_node;
}

void del(int nodeId){
	int child = 0;
	if(lC[nodeId] != -1){ child += 1; }
	if(rC[nodeId] != -1){ child += 1; }
	int p;
	int c;
	if(child == 0){
		p = parent[nodeId];
		if(p == -1){
			root = -1;
		}else if(lC[p] == nodeId){
			lC[p] = -1;
		}else{
			rC[p] = -1;
		}
		
	}else if(child == 1){
		if(lC[nodeId] != -1){
			c = lC[nodeId];
		}else{
			c = rC[nodeId];
		}
		p = parent[nodeId];
		if(p == -1){
			root = c;
		}else{
			if(lC[p] == nodeId){
				lC[p] = c;
			}else{
				rC[p] = c;
			}
		} 
		parent[c] = p;
	}else{
		int rep = find_min(rC[nodeId]);
		val[nodeId] = val[rep];
		del(rep);
	}
}

int gen(int start, int end){
	if (start > end){
		return -1;
	}
	int mid = (start+end)/2;
	int nodeVal = openings[mid];
	int nodeId = val.size();
	int temp;
	val.push_back(nodeVal);
	lC.push_back(-1);
	rC.push_back(-1);
	parent.push_back(-1);
	
	temp = gen(start, mid-1);
	lC[nodeId] = temp;
	
	temp = gen(mid+1, end); 
	rC[nodeId] = temp;
	
	if(lC[nodeId] != -1){
		parent[lC[nodeId]] = nodeId;
	}
	if(rC[nodeId] != -1){
		parent[rC[nodeId]] = nodeId;
	}
	
	return nodeId;
}

int main(){	
	int n, k, m;
	char e;
	int c;
	int o = 0;
	bool possible = true;
	cin >> n >> k >> m;
	vector<int> catches(0); //sorted list of catches
	vector<bool> eventType(0); //true - catch, false openings
	vector<int> eventCarrier(0); //who did that
	vector<int> eventPrison(0); //where it happened
	
	vector<int> cE(n, -1); //thiev -> eventId when caught if unclaimed, else -1;
	
	for(int i = 0; i < m; i++){ // i is current event
		cin >> e >> c;
		if(e == 'C'){
			eventType.push_back(true);
			eventCarrier.push_back(c-1);
			eventPrison.push_back(-1);
			
		}else{
			eventType.push_back(false);
			eventCarrier.push_back(c-1);
			eventPrison.push_back(o);
			o = (o+1)%k;
			openings.push_back(i);
		}
	}
	
	root = gen(0, openings.size() - 1);
	bool type;
	int carrier;
	int bN;
	
	for(int i = 0; i < m; i++){
		type = eventType[i];
		carrier = eventCarrier[i];
		if(cE[carrier] != -1){ //gotta get my homie out
			bN = query(cE[carrier]);
			//validate opening
			auto minO = lower_bound(openings.begin(), openings.end(), cE[carrier]);
			auto maxO = minO;
			for(int i = 0; i < k-1; i++){
				maxO++;
				if(maxO == openings.end()){
					maxO--;
					break;
				}
			}
			if(bN == -1 or val[bN] > *maxO){ //ugyldig
				if(minO == openings.end() or *minO > i){ //homie rekker ikke å komme ut D:
					possible = false;
					break;
				}else{
					eventPrison[cE[carrier]] = eventPrison[*minO];
					cE[carrier] = -1;
				}
			}else{ //gyldig
				eventPrison[cE[carrier]] = eventPrison[val[bN]];
				cE[carrier] = -1;
				del(bN);
			}
		}
		if(type == true){
			cE[carrier] = i;
		}
	}
	
	for(int i = 0; i < n; i++){
		if(cE[i] != -1){
			bN = query(cE[i]);
			auto minO = lower_bound(openings.begin(), openings.end(), cE[i]);
			auto maxO = minO;
			for(int i = 0; i < k-1; i++){
				maxO++;
				if(maxO == openings.end()){
					maxO--;
					break;
				}
			}
			if(bN == -1 or val[bN] > *maxO){
				eventPrison[cE[i]] = 0;
				cE[i] = -1;
			}else{
				eventPrison[cE[i]] = eventPrison[val[bN]];
				cE[i] = -1;
				del(bN);
			}
		}
	}
	
	if(root != -1){
		possible = false;
	}
	
	if(not possible){
		cout << "IMPOSSIBLE" << endl;
	}else{
		for(int i = 0; i < m; i++){
			cout << eventPrison[i] + 1 << " ";
		}
		cout << endl;
	}
	
	return 0;	
}

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: WRONG ANSWER

Test 5

Group: 2, 4, 5

Verdict: ACCEPTED

Test 6

Group: 4, 5

Verdict: ACCEPTED

Test 7

Group: 4, 5

Verdict: WRONG ANSWER

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: WRONG ANSWER

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: TIME LIMIT EXCEEDED

Test 20

Group: 2, 5

Verdict: TIME LIMIT EXCEEDED

Test 21

Group: 2, 5

Verdict: TIME LIMIT EXCEEDED

Test 22

Group: 5

Verdict: ACCEPTED

Test 23

Group: 5

Verdict: ACCEPTED

Test 24

Group: 3, 4, 5

Verdict: WRONG ANSWER

Test 25

Group: 3, 4, 5

Verdict: WRONG ANSWER

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: WRONG ANSWER

Test 37

Group: 3, 5

Verdict: WRONG ANSWER

Test 38

Group: 3, 5

Verdict: WRONG ANSWER

Test 39

Group: 3, 5

Verdict: WRONG ANSWER

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