Thieves and Prisons

#include <bits/stdc++.h>
#define pb push_back
#define fi first
#define se second

using namespace std;

typedef pair<int, int> pii;
typedef vector<int> vi;
typedef vector<pii> vpii;

const int N = 100005;

vpii event[N]; //second true if he has been caught
bool sol = true;
int ans[N];
vi open;

struct seg{
    int a, b, id;
    bool operator < (const seg &o) const{
        return b > o.b; 
    }
};

bool cmp(const seg &lhs, const seg &rhs){
    return lhs.a < rhs.a;
}

void display(const seg &v){
    cout << v.a << " " << v.b << " " << v.id << endl;
}

vector<seg> vec;
priority_queue<seg> pq;

int main(){
    int n, k, m, x; char t;
    cin >> n >> k >> m;
    for(int i = 0; i < m; i++){
        cin >> t >> x;
        if(t == 'C') event[x].pb({i, true}); 
        else{
            event[x].pb({i, false});
            open.pb(i);
            ans[i] = open.size(); //TO CHANGE
        }
    }

    for(int i = 1; i <= n; i++){
        for(int j = 0; j < int(event[i].size()) - 1; j++)
            if(event[i][j].se)
                vec.pb({event[i][j].fi + 1, event[i][j + 1].fi - 1, i});
        if(!event[i].empty() && event[i].back().se) vec.pb({event[i].back().fi + 1, m + 1, i});
    }

    sort(vec.begin(), vec.end(), cmp);


    //for(auto u : vec) display(u);

    int idx = 0;
    for(int i = 0; i < int(open.size()) && sol; i++){
        //cout << i << ", " << open[i] << ": " << endl;
        while(idx < int(vec.size()) && vec[idx].a <= open[i]) pq.push(vec[idx++]);
        int cnt = 0;

        //cout << "top: ";
        //display(pq.top());
        //cout << "size " << pq.size() << endl;

        if(pq.empty() || pq.top().b < open[i]){sol = false; break;}

        while(!pq.empty() && (i == int(open.size()) - 1 || pq.top().b < open[i + 1])){
            cnt++;
            ans[pq.top().a - 1] = i + 1;
            pq.pop();
        }
        if(!sol || (cnt == 0 && pq.empty())){sol = false; break;}

        //cout << ": " << cnt << endl;
        //cout << "size " << pq.size() << endl;

        if(cnt == 0){
            ans[pq.top().a - 1] = i + 1;
            pq.pop();
        }

    }
    
    while(!pq.empty()){
        if(pq.top().b != m + 1) sol = false;
        ans[pq.top().a - 1] = 1;
        pq.pop();
    }
    for(int i = idx; i < int(vec.size()); i++){
        if(vec[i].b != m + 1) sol = false;
        ans[vec[i].a - 1] = 1;
    }

    if(!sol){
        cout << "IMPOSSIBLE\n";
        return 0;
    }

    for(int i = 0; i < m; i++)
        cout << ans[i] << " ";
    cout << endl;
}

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

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

Test 27

Group: 3, 4, 5

Verdict: WRONG ANSWER

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