Falling Squares LeetCode solution in Python

1
class SegmentTreeNode:
2
    def __init__(self, low, high):
3
        self.low = low
4
        self.high = high
5
        self.left = None
6
        self.right = None
7
        self.max = 0
8

9

10
class Solution:
11
    def _build(self, left, right):
12
        root = SegmentTreeNode(self.coords[left], self.coords[right])
13
        if left == right:
14
            return root
15

16
        mid = (left + right) // 2
17
        root.left = self._build(left, mid)
18
        root.right = self._build(mid + 1, right)
19
        return root
20

21
    def _update(self, root, lower, upper, val):
22
        if not root:
23
            return
24
        if lower <= root.high and root.low <= upper:  # intersect
25
            root.max = val
26
            self._update(root.left, lower, upper, val)
27
            self._update(root.right, lower, upper, val)
28

29
    def _query(self, root, lower, upper):
30
        if lower <= root.low and root.high <= upper:
31
            return root.max
32
        if upper < root.low or root.high < lower:
33
            return 0
34
        return max(
35
            self._query(root.left, lower, upper), self._query(root.right, lower, upper)
36
        )
37

38
    def fallingSquares(self, positions: List[List[int]]) -> List[int]:
39
        # coordinates compression
40
        coords = set()
41
        for left, size in positions:
42
            right = left + size - 1
43
            coords.add(left)
44
            coords.add(right)
45
        self.coords = sorted(list(coords))
46
        root = self._build(0, len(self.coords) - 1)
47

48
        res = []
49
        for left, size in positions:
50
            right = left + size - 1
51
            h = self._query(root, left, right) + size
52
            res.append(max(res[-1], h)) if res else res.append(h)
53
            self._update(root, left, right, h)
54
        return res