Exam Room Python Solution with Tests

lcpy gen -n 855   # by problem number
lcpy gen -s exam_room   # by problem name

Input
["ExamRoom", "seat", "seat", "seat", "seat", "leave", "seat"]
[[10], [], [], [], [], [4], []]
Output
[null, 0, 9, 4, 2, null, 5]

Explanation
ExamRoom examRoom = new ExamRoom(10);
examRoom.seat(); // return 0, no one is in the room, then the student sits at seat number 0.
examRoom.seat(); // return 9, the student sits at the last seat number 9.
examRoom.seat(); // return 4, the student sits at the last seat number 4.
examRoom.seat(); // return 2, the student sits at the last seat number 2.
examRoom.leave(4);
examRoom.seat(); // return 5, the student sits at the last seat number 5.

import heapq
from bisect import bisect_left, insort


class ExamRoom:
    # Time: __init__ O(1), seat O(log k) amortized, leave O(log k + k)
    # Space: O(k) for k occupied seats

    def __init__(self, n: int) -> None:
        self.n = n
        self.occupied: list[int] = []
        self.gaps: list[tuple[int, int, int]] = []
        self._push(-1, n)

    def _push(self, left: int, right: int) -> None:
        if right - left < 2:
            return
        if left == -1:
            dist = right
        elif right == self.n:
            dist = self.n - 1 - left
        else:
            dist = (right - left) // 2
        heapq.heappush(self.gaps, (-dist, left, right))

    def _valid(self, left: int, right: int) -> bool:
        occ = self.occupied
        if left == -1:
            return (right == self.n) if not occ else occ[0] == right
        if right == self.n:
            return occ[-1] == left
        i = bisect_left(occ, left)
        return i + 1 < len(occ) and occ[i] == left and occ[i + 1] == right

    def seat(self) -> int:
        while True:
            _, left, right = heapq.heappop(self.gaps)
            if not self._valid(left, right):
                continue
            if left == -1:
                pos = 0
            elif right == self.n:
                pos = self.n - 1
            else:
                pos = (left + right) // 2
            insort(self.occupied, pos)
            self._push(left, pos)
            self._push(pos, right)
            return pos

    def leave(self, p: int) -> None:
        self.occupied.remove(p)
        i = bisect_left(self.occupied, p)
        left = self.occupied[i - 1] if i > 0 else -1
        right = self.occupied[i] if i < len(self.occupied) else self.n
        self._push(left, right)