Exam Room LeetCode solution in Python

1
# Runtime: 182 ms (Top 65.24%) | Memory: 20.7 MB (Top 5.08%)
2
class ExamRoom:
3

4
    def __init__(self, n: int):
5
        self.N = n
6
        self.pq = []
7
        self.dict = {}
8
        self.addSegment(0, self.N - 1)
9

10
    def seat(self) -> int:
11
        start, end, distance = heapq.heappop(self.pq)
12
        self.dict.pop(start, None)  # Remove old segment from dictionary
13
        self.dict.pop(end, None)
14

15
        if start == end:
16
            position = start
17

18
        elif start == 0:
19
            position = start
20
            right = self.addSegment(start + 1, end)
21

22
        elif end == self.N - 1:
23
            position = end
24
            left = self.addSegment(start, end - 1)
25

26
        elif end - start == 1:  # ONLY ONE PLACE TO PUT
27
            position = start
28
            left = self.addSegment(start + 1, end)
29

30
        else:
31
            position = start + (end - start) // 2
32
            right = self.addSegment(start, position - 1)
33
            left = self.addSegment(position + 1, end)
34

35
        return position
36

37
    def leave(self, p: int) -> None:
38
        left = self.dict.get(p - 1, None)
39
        right = self.dict.get(p + 1, None)
40

41
        new_start = new_end = p
42

43
        if left:
44
            self.removeSegment(left)
45
            new_start = left.start
46

47
        if right:
48
            self.removeSegment(right)
49
            new_end = right.end
50

51
        self.addSegment(new_start, new_end)
52

53
    def addSegment(self, start, end):
54
        segment = Segment(start, end, self.N)
55
        self.dict[segment.start] = segment
56
        self.dict[segment.end] = segment
57
        heapq.heappush(self.pq, segment)
58

59
    def removeSegment(self, segment):
60
        self.dict.pop(segment.start, None)
61
        self.dict.pop(segment.end, None)
62
        self.pq.remove(segment)
63

64

65
class Segment:
66
    def __init__(self, start, end, N):
67
        self.start = start
68
        self.end = end
69
        self.distance = self.calculateDistance(start, end, N)
70

71
    def __lt__(self, other_segment):
72
        return (
73
            self.distance > other_segment.distance
74
            if self.distance != other_segment.distance
75
            else self.start < other_segment.start
76
        )
77

78
    def calculateDistance(self, start, end, N):
79
        if start == 0:
80
            return end
81

82
        if end == N - 1:
83
            return end - start
84

85
        else:
86
            return (end - start) // 2
87

88
    def __iter__(self):
89
        return iter((self.start, self.end, self.distance))