Leetcode 831 Solution

This article provides solution to leetcode question 831 (largest-sum-of-averages)

https://leetcode.com/problems/largest-sum-of-averages

Solution

class Solution:
    def largestSumOfAverages(self, A: List[int], K: int) -> float:
        self.memo = {}
        self.presum = [0]
        for a in A:
            self.presum.append(self.presum[-1] + a)

        def dfs(A, i, K):
            if (i, K) in self.memo:
                return self.memo[(i, K)]

            if i == len(A):
                return 0 if K == 0 else -1

            if K == 1:
                return (self.presum[len(A)] - self.presum[i]) / (len(A) - i)

            opt = -1
            for j in range(i, len(A)):
                s = dfs(A, j + 1, K - 1)

                if s != -1:
                    opt = max(opt, s + (self.presum[j + 1] - self.presum[i]) / (j - i + 1))

            self.memo[(i, K)] = opt
            return opt

        return dfs(A, 0, K)