Incremental Memory Leak LeetCode solution in Python

1
class Solution:
2
    def memLeak(self, memory1: int, memory2: int) -> List[int]:
3

4
        inverted = False
5
        if memory2 > memory1:
6
            memory2, memory1 = memory1, memory2
7
            inverted = True
8

9
            # Compute the number of steps in first stage - 1
10
        i_start = solve_quadratic(1, 2 * (memory1 - memory2))
11

12
        # Memory1 after the end of first stage is computed using the sum of arithmetic sequence
13
        memory1 -= i_start * (i_start + 1) // 2
14

15
        if (
16
            memory1 == memory2
17
        ):  # Special case (if we end up with equal numbers after stage - 1 - undo inversion)
18
            inverted = False
19

20
        # Compute number of steps in stage - 2
21
        n_end = solve_quadratic((i_start + 1), memory2)
22

23
        # Compute sums of respective arithmetic sequences
24
        i_end_1 = i_start - 1 + 2 * n_end
25
        i_end_2 = i_start + 2 * n_end
26

27
        sum1 = n_end * (i_start + 1 + i_end_1) // 2
28
        sum2 = n_end * (i_start + 2 + i_end_2) // 2
29

30
        # Compute updated memories
31
        memory1 -= sum1
32
        memory2 -= sum2
33

34
        full_cnt = 2 * n_end + i_start
35

36
        if (
37
            memory1 >= i_end_2 + 1
38
        ):  # If we can still make one removal from the first stick - perform it.
39
            memory1 -= i_end_2 + 1
40
            full_cnt += 1
41

42
        return (
43
            [full_cnt + 1, memory2, memory1]
44
            if inverted
45
            else [full_cnt + 1, memory1, memory2]
46
        )