dp
dfs
accumulate
]
Leetcode 0813 Largest Sum of Averages
Problem statement
https://leetcode.com/problems/largest-sum-of-averages/
Solution
Let dfs(i, k)
be answer for A[:i+1]
with we need to split into k
parts. Then, first of all, we consider only cases, when i >= k-1
, and if i == k - 1
, we return cumulative sum of first i
elements: it means, that all groups should have length $1$. Also, if we have only 1
group, that is k = 1
, we return CS[i]/(i+1)
. Finally, we do recursion: for j in range(k - 2, i)
, we run dfs(j, k - 1)
and evaluate mean in last group.
Complexity
Time complexity is O(n^2k)
: we have nk
different states and at most k
transitions from one state to another. Space complexity is O(nk)
, it can be reduced to O(n)
, because each level we increase $k$ only by one.
Code
class Solution:
def largestSumOfAverages(self, A, K):
CS = list(accumulate(A))
@lru_cache(None)
def dfs(i, k):
if k == 1: return CS[i]/(i+1)
if i == k - 1: return CS[i]
return max(dfs(j, k - 1) + (CS[i] - CS[j])/(i-j) for j in range(k-2, i))
return dfs(len(A) - 1, K)
Remark
I wonder if it is possible to do something like Knuth optimization here with total complexity O(nk)
.