Maximum Sum of 3 Non-Overlapping Subarrays LeetCode solution in Python

1
from itertools import accumulate
2
from functools import lru_cache
3

4

5
class Solution:
6
    def maxSumOfThreeSubarrays(self, nums: List[int], k: int) -> List[int]:
7
        n = len(nums)
8
        windows = list(accumulate(nums))
9
        windows = [
10
            windows[i + k - 1] - (windows[i - 1] if i > 0 else 0)
11
            for i in range(len(windows) - k + 1)
12
        ]
13

14
        @lru_cache(None)
15
        def dfs(i, t):
16
            if t == 0:
17
                return 0, []
18
            if i >= len(windows):
19
                return float("-inf"), []
20
            cost1, sol1 = dfs(i + k, t - 1)
21
            cost2, sol2 = dfs(i + 1, t)
22
            if windows[i] + cost1 < cost2:
23
                return cost2, sol2
24
            return windows[i] + cost1, [i] + sol1
25

26
        return dfs(0, 3)[1]