Booking Concert Tickets in Groups LeetCode solution in Python

1
class Node:
2
    def __init__(self, start, end):
3
        self.s = start
4
        self.e = end
5
        self.left = None
6
        self.right = None
7
        self.total = 0  # for range sum query
8
        self.mx = 0  # for range max query
9

10

11
class SegTree:
12
    def __init__(self, start, end, val):
13

14
        def build(l, r):
15
            if l > r:
16
                return None
17
            if l == r:
18
                node = Node(l, r)
19
                node.total = val
20
                node.mx = val
21
                return node
22
            node = Node(l, r)
23
            m = (l + r) // 2
24
            node.left = build(l, m)
25
            node.right = build(m + 1, r)
26
            node.mx = max(node.left.mx, node.right.mx)
27
            node.total = node.left.total + node.right.total
28
            return node
29

30
        self.root = build(start, end)
31

32
    # update the total remain seats and the max remain seats for each node (range) in the segment tree
33
    def update(self, index, val):
34

35
        def updateHelper(node):
36
            if node.s == node.e == index:
37
                node.total -= val
38
                node.mx -= val
39
                return
40
            m = (node.s + node.e) // 2
41
            if index <= m:
42
                updateHelper(node.left)
43
            elif index > m:
44
                updateHelper(node.right)
45
            node.mx = max(node.left.mx, node.right.mx)
46
            node.total = node.left.total + node.right.total
47
            return
48

49
        updateHelper(self.root)
50

51
    def maxQuery(self, k, maxRow, seats):
52

53
        def queryHelper(node):
54
            if node.s == node.e:
55
                # check if the row number is less than maxRow and the number of remains seats is greater or equal than k
56
                if node.e > maxRow or node.total < k:
57
                    return []
58
                if node.e <= maxRow and node.total >= k:
59
                    return [node.e, seats - node.total]
60
            # we want to greedily search the left subtree to get the smallest row which has enough remain seats
61
            if node.left.mx >= k:
62
                return queryHelper(node.left)
63
            return queryHelper(node.right)
64

65
        return queryHelper(self.root)
66

67
    def sumQuery(self, endRow):
68

69
        def queryHelper(node, left, right):
70
            if left <= node.s and node.e <= right:
71
                return node.total
72
            m = (node.s + node.e) // 2
73
            if right <= m:
74
                return queryHelper(node.left, left, right)
75
            elif left > m:
76
                return queryHelper(node.right, left, right)
77
            return queryHelper(node.left, left, m) + queryHelper(
78
                node.right, m + 1, right
79
            )
80

81
        return queryHelper(self.root, 0, endRow)
82

83

84
class BookMyShow:
85

86
    def __init__(self, n: int, m: int):
87
        self.m = m
88
        self.seg = SegTree(0, n - 1, m)
89
        # record the remain seats at each row
90
        self.seats = [m] * n
91
        # record the index of the smallest row that has remain seats > 0
92
        self.startRow = 0
93

94
    def gather(self, k: int, maxRow: int) -> List[int]:
95
        res = self.seg.maxQuery(k, maxRow, self.m)
96
        if res:
97
            row = res[0]
98
            self.seg.update(row, k)
99
            self.seats[row] -= k
100
        return res
101

102
    def scatter(self, k: int, maxRow: int) -> bool:
103
        if self.seg.sumQuery(maxRow) < k:
104
            return False
105
        else:
106
            i = self.startRow
107
            total = 0
108
            while total < k:
109
                prevTotal = total
110
                total += self.seats[i]
111
                if total < k:
112
                    # use up all the seats at ith row
113
                    self.seg.update(i, self.seats[i])
114
                    self.seats[i] = 0
115
                    i += 1
116
                    self.startRow = i
117
                elif total >= k:
118
                    # occupy (k - prevTotal) seats at ith row
119
                    self.seg.update(i, k - prevTotal)
120
                    self.seats[i] -= k - prevTotal
121
            return True