To Solve this problem, we can simply use Dynamic Programming.
We can get solution(A[:i]) based on solution(A[:j]) where i - K <= j < i. So the time complexity is O(K * len(A)). As the limitaion of input is 1 <= K <= A.length <= 500 0 <= A[i] <= 10^6, the overall calculatations would be around 5*10^8 which would be alright.
One problem here is that, the original tags of this problem on leetcode is graph. I can't get why and how to solve this problem with graph.
class Solution:
def maxSumAfterPartitioning(self, A, K):
dp = [0] * (len(A) + 1)
for end in range(1, len(A) + 1):
m = A[end - 1]
for last_part in range(1, K + 1):
last_end = end - last_part
if last_end < 0:
break
m = max(A[last_end], m)
dp[end] = max(dp[end], dp[last_end] + last_part * m)
return dp[-1]