leetcode/solution/0800-0899/0855.Exam Room/Solution.py

46 lines
1.1 KiB
Python

class ExamRoom:
def __init__(self, n: int):
def dist(x):
l, r = x
return r - l - 1 if l == -1 or r == n else (r - l) >> 1
self.n = n
self.ts = SortedList(key=lambda x: (-dist(x), x[0]))
self.left = {}
self.right = {}
self.add((-1, n))
def seat(self) -> int:
s = self.ts[0]
p = (s[0] + s[1]) >> 1
if s[0] == -1:
p = 0
elif s[1] == self.n:
p = self.n - 1
self.delete(s)
self.add((s[0], p))
self.add((p, s[1]))
return p
def leave(self, p: int) -> None:
l, r = self.left[p], self.right[p]
self.delete((l, p))
self.delete((p, r))
self.add((l, r))
def add(self, s):
self.ts.add(s)
self.left[s[1]] = s[0]
self.right[s[0]] = s[1]
def delete(self, s):
self.ts.remove(s)
self.left.pop(s[1])
self.right.pop(s[0])
# Your ExamRoom object will be instantiated and called as such:
# obj = ExamRoom(n)
# param_1 = obj.seat()
# obj.leave(p)