Untitled

#include <iostream>
#include <set>
#include <unordered_map>
#include <vector>

using namespace std;

struct cmp {
    bool operator()(const pair<int, int>& a, const pair<int, int>& b) const {
        if ((a.second - a.first) == (b.second - b.first)) {
            return a.first < b.first;
        }
        return (a.second - a.first) > (b.second - b.first);
    }
};

set<pair<int, int>, cmp> space;
set<int> locker;
unordered_map<int, int> mp;
int mid;
int n;
int remain;

void init(int N) {
    n = N;
    remain = N;
    mid = -1;
    mp.clear();
    space.clear();

    space.insert(make_pair(1, N));
}

int arrive(int mId) {
    remain--;
    auto it = space.begin();
    auto st = it->first;
    auto en = it->second;

    if (st == 1) {
        mid = 1;
        space.insert(make_pair(st + 1, n));
    } else if (en == n) {
        mid = n;
        space.insert(make_pair(st, n - 1));
    } else {
        mid = (st + en) / 2;
        space.insert(make_pair(st + 1, mid - 1));
        space.insert(make_pair(mid + 1, en - 1));
    }

    locker.insert(mid);
    mp[mId] = mid;
    return mid;
}

int leave(int mId) {
    remain++;
    int idx = mp[mId];
    auto it = locker.find(idx);
    auto nex = next(it);
    auto pre = prev(it);

    auto tmp1 = make_pair(*pre + 1, *it - 1);
    auto tmp2 = make_pair(*it + 1, *nex - 1);

    space.erase(tmp1);
    space.erase(tmp2);

    mp.erase(mId);

    locker.erase(it);

    return remain;
}