Grumpy Bookstore Owner LeetCode solution in Python

1
class Solution:
2
    def recursion(self, index, used):
3
        # base case
4
        if index == self.n:
5
            return 0
6

7
        # check in dp
8
        if (index, used) in self.dp:
9
            return self.dp[(index, used)]
10
        # choice1 is using the secret technique
11
        choice1 = -float("inf")
12

13
        # we can only use secret technique once and consecutively
14
        if used == True:
15
            # use the secret technique
16
            end = index + self.minutes if index + self.minutes < self.n else self.n
17
            to_substract = self.prefix_sum[index - 1] if index != 0 else 0
18
            val = self.prefix_sum[end - 1] - to_substract
19
            choice1 = self.recursion(end, False) + val
20

21
        # Do not use the secret tehcnique and play simple
22
        choice2 = self.recursion(index + 1, used) + (
23
            self.customers[index] if self.grumpy[index] == 0 else 0
24
        )
25
        ans = choice1 if choice1 > choice2 else choice2
26

27
        # Memoization is done here
28
        self.dp[(index, used)] = ans
29
        return ans
30

31
    def maxSatisfied(
32
        self, customers: List[int], grumpy: List[int], minutes: int
33
    ) -> int:
34
        self.n = len(customers)
35
        self.customers = customers
36
        self.grumpy = grumpy
37
        self.minutes = minutes
38
        self.dp = {}
39
        self.prefix_sum = [x for x in customers]
40
        for i in range(1, self.n):
41
            self.prefix_sum[i] += self.prefix_sum[i - 1]
42
        return self.recursion(0, True)