Problem statement

https://leetcode.com/problems/exam-room/

Solution

There is straightforward solution with O(n) complexity to seat and leave, we just keep sorted list and find shortest distance. However we can do better: let us have:

  1. self.places are occupied places, we keep them in SortedList python object
  2. self.scores are segments between two adjacent persons with scores: imagine, that we have people on places [-1, 4, 9, 12], where n = 12 (we use two dummy seats: -1 and n). Then we have 3 possible segments where we can seat people: [-1, 4] and score here is 4, because putting person on seat 0 it will be 4 seats away from any other one. Next segment is [4, 9] and we have score 2 here, if we put new person on seat 6. Finally, for segment [9, 12] score is 2, because we can put person on seat number 11. In the beginning we create one segment [-1, n] and its score.
  3. priority function will evaluate score of segment, here we need to deal with two cases: when we are in the middle or not.
  4. seat function will work like this: first, we found the biggest element from self.scores, it will be the segment we need to seat person to. Then, if left end of this segment is -1, we need to seat person on place 0, if right end is n, we need to seat person on place n - 1, in other case we seat him in the middle. We also add two new segments to self.scores, one place to self.places and return mid.
  5. leave Let us find index of person we need to delete, find its neighbours, delete one segment and add two new segments from to self.scores and remove place from self.places.

Complexity

Time complexity of both operations is O(log n), where n is current number of occupied seats.

Code

from sortedcontainers import SortedList

class ExamRoom:
    def __init__(self, N):
        self.N = N
        self.places = SortedList([-1, N])
        self.scores = SortedList()
        self.priority(-1,N,SortedList.add)
        
    def priority(self, beg, end, func): 
        if beg == -1 or end == self.N:
            score = beg - end + 1
        else:
            score = -((end - beg) // 2)
        func(self.scores, (score, beg, end))
        
    def seat(self):
        _, beg, end  = self.scores.pop(0)
        
        if beg == -1:
            mid = beg + 1
        elif end == self.N:
            mid = end - 1
        else:
            mid = (beg + end) // 2

        self.priority(beg, mid, SortedList.add)
        self.priority(mid, end, SortedList.add)
    
        self.places.add(mid)
        return mid

    def leave(self, p):
        ind = self.places.index(p)
        beg, end = self.places[ind-1], self.places[ind+1]
        
        self.priority(beg, end, SortedList.add)
        self.priority(beg, p, SortedList.remove)
        self.priority(p, end, SortedList.remove)
        
        self.places.remove(p)

Remark

There is alternative solution with the same complexity, using heaps with lazy updates. We need to keep tuples of four elements: last one is flag if this segment deleted or not. Also we need to keep dictionary (or two dictionaries), which connects each point to left and right segments.