Thieves and Prisons

#include <bits/stdc++.h>
using namespace std;

template <typename T> T load() { T r; cin >> r; return r; }
template <typename T> vector<T> loadMany(int n) { vector<T> rs(n); generate(rs.begin(), rs.end(), &load<T>); return rs; }

template <typename T> typename T::value_type extract(typename T::iterator it, T& xs) { auto x = *it; xs.erase(it); return x; }

struct EscapeReq {
    int thief, t1, t2;
    int source() { return t1-1; }
    static bool compT1(EscapeReq a, EscapeReq b) { return a.t1 < b.t1; }
    static bool compT2(EscapeReq a, EscapeReq b) { return a.t2 < b.t2; }
};
struct Event {
    char type;
    int thief;
    friend istream& operator>>(istream& in, Event& ev) { in >> ev.type >> ev.thief; --ev.thief; return in; }
};

struct FluxPrison {
    int lastBlock, nextBlock, prison, lastOccIndex;
    static FluxPrison id() { return FluxPrison{-1, numeric_limits<int>::max(), -1, -1}; }
    static FluxPrison merge(FluxPrison a, FluxPrison b) {
        return min(a, b, [&](FluxPrison c, FluxPrison d){
            return c.nextBlock < d.nextBlock;
        });
    }
};

template <typename T> struct Segtree {
    Segtree(int n):base(1){
        while (base < n) base *= 2;
        nodes.resize(base*2, T::id());
    }
    void setLeaf(int i, T x) {
        nodes[base+i] = move(x);
        for (auto v=(base+i)/2; v>0; v/=2)
            nodes[v] = T::merge(nodes[v*2], nodes[v*2+1]);
    }
    T readRange(int left, int right) const { return implRR(left, right, 0, base, 1); }
    T implRR(int left, int right, int i, int e, int v) const {
        if (left >= e or right <= i) return T::id();
        if (left <= i and right >= e) return nodes[v];
        return T::merge(implRR(left, right, i, (i+e)/2, v*2), implRR(left, right, (i+e)/2, e, v*2+1));
    }
    int base;
    vector<T> nodes;
};

vector<int> solve(int n, int k, int m, const vector<Event>& events) {
    auto occs = vector<vector<int>>(n);
    for (auto i=0; i<m; ++i)
        occs[events[i].thief].push_back(i);
    
    auto reqs = vector<EscapeReq>();
    for (auto i=0; i<n; ++i)
        for (auto j=0; j+1<(int)occs[i].size(); ++j)
            if (events[occs[i][j]].type == 'C')
                reqs.push_back(EscapeReq{i, occs[i][j]+1, occs[i][j+1]});

    sort(reqs.begin(), reqs.end(), &EscapeReq::compT1);

    // phase 1
    //   - no assignments are made
    //   - make sure thieves required to be free are free

    auto qr = set<EscapeReq, decltype(&EscapeReq::compT2)>(&EscapeReq::compT2);
    auto reqIt = reqs.begin();
    auto assignment = vector<int>(m, -2);
    auto lastPrison = -1;
    auto lastOpening = -1;
    for (auto i=0; i<m; ++i) {
        for (;reqIt!=reqs.end() and reqIt->t1 == i; ++reqIt)
            qr.insert(*reqIt);
        while (not qr.empty() and qr.begin()->t2 <= i) {
            auto req = *qr.begin();
            qr.erase(qr.begin());
            if (lastOpening >= req.t1)
                assignment[req.source()] = lastPrison;
            else
                return {};
        }
        if (events[i].type == 'O') {
            lastOpening = i;
            if (not qr.empty()) {
                auto req = *qr.begin();
                qr.erase(qr.begin());
                assignment[i] = assignment[req.source()] = lastPrison = (lastPrison + 1) % k; 
            }
        }
    }

    // phase 2
    //   - all thieves who were required to be free are free
    //   - make sure no escapes without closing were made

    auto filledPrisons = set<int>();
    auto emptyPrisons = set<int>();
    auto blocks = vector<vector<int>>(k);
    auto prevblock = vector<int>(m, -1);
    for (auto i=0; i<k; ++i)
        emptyPrisons.insert(i);
    for (auto i=0; i<m; ++i) {
        if (events[i].type == 'C') {
            if (assignment[i] != -2) {
                if (emptyPrisons.count(assignment[i]))
                    emptyPrisons.erase(emptyPrisons.find(assignment[i])), filledPrisons.insert(assignment[i]);
            }
        } else {
            if (assignment[i] != -2) {
                if (not blocks[assignment[i]].empty())
                    prevblock[i] = blocks[assignment[i]].back();
                else
                    prevblock[i] = -1 - assignment[i];
                blocks[assignment[i]].push_back(i);
                if (not filledPrisons.count(assignment[i]))
                    return {};
                else
                    filledPrisons.erase(filledPrisons.find(assignment[i])), emptyPrisons.insert(assignment[i]);
            }
        }
    }

    // phase 3
    //   - everything so far makes sense and any matching will be correct
    //   - assign unassigned things

    auto closes = set<int>();
    auto prisons = Segtree<FluxPrison>(k+m+k);
    auto pseudoNextBlock = [&](int prison, int index){
        return index < (int)blocks[prison].size() ? blocks[prison][index] : m+prison;
    };
    for (auto i=0; i<k; ++i) {
        auto nb = pseudoNextBlock(i, 0);
        prisons.setLeaf(k+(-1-i), FluxPrison{-1-i, nb, i, 0});
    }

    for (auto i=0; i<m; ++i) {
        if (events[i].type == 'O') {
            if (assignment[i] != -2) {
                auto id = assignment[i];
                auto b1 = prevblock[i];
                auto b2 = i;
                auto old = prisons.readRange(b1+k, b1+k+1);
                auto b3 = pseudoNextBlock(id, old.lastOccIndex+1);
                prisons.setLeaf(b1+k, FluxPrison::id());
                prisons.setLeaf(b2+k, FluxPrison{b2, b3, id, old.lastOccIndex+1});
            } else {
                if (closes.empty())
                    return {};
                auto maxfc = *prev(closes.end());
                auto prison = prisons.readRange(0, k+maxfc+1);
                if (prison.prison == -1)
                    return {};
                auto close = extract(closes.lower_bound(prison.lastBlock), closes);
                assignment[close] = assignment[i] = prison.prison;
                auto b1 = prison.lastBlock;
                auto b2 = i;
                auto b3 = prison.nextBlock;
                prisons.setLeaf(b1+k, FluxPrison::id());
                prisons.setLeaf(b2+k, FluxPrison{b2, b3, prison.prison, prison.lastOccIndex});
            }
        } else {
            if (assignment[i] == -2) {
                closes.insert(i);
            }
        }
    }

    // phase 4
    //   - all openings are correctly assigned
    //   - assigned closing anyhow
    for (auto i=0; i<m; ++i)
        if (events[i].type == 'C' and assignment[i] == -2)
            assignment[i] = 0;

    return assignment;
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    auto n = load<int>(); // thief count
    auto k = load<int>(); // prison count
    auto m = load<int>(); // event count
    auto events = loadMany<Event>(m);
    auto answer = solve(n, k, m, events);
    if (not answer.empty()) {
        for (auto i : answer)
            cout << i+1 << ' ';
        cout << '\n';
    } else
        cout << "IMPOSSIBLE\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: 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: WRONG ANSWER

Test 29

Group: 4, 5

Verdict: WRONG ANSWER

Test 30

Group: 4, 5

Verdict: WRONG ANSWER

Test 31

Group: 4, 5

Verdict: WRONG ANSWER

Test 32

Group: 2, 4, 5

Verdict: WRONG ANSWER

Test 33

Group: 2, 4, 5

Verdict: WRONG ANSWER

Test 34

Group: 2, 4, 5

Verdict: WRONG ANSWER

Test 35

Group: 2, 4, 5

Verdict: WRONG ANSWER

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

Test 41

Group: 5

Verdict: WRONG ANSWER

Test 42

Group: 5

Verdict: WRONG ANSWER

Test 43

Group: 5

Verdict: WRONG ANSWER

Test 44

Group: 2, 5

Verdict: WRONG ANSWER

Test 45

Group: 2, 5

Verdict: WRONG ANSWER

Test 46

Group: 2, 5

Verdict: WRONG ANSWER

Test 47

Group: 2, 5

Verdict: WRONG ANSWER