Problem statement

https://leetcode.com/problems/maximum-subarray-sum-after-one-operation/

Solution

The idea is to use dp with state dp(i, j), where `

  1. dp(i, 1) is answer to subproblem nums[:i+1] where we did not use square yet.
  2. dp(i, 0) is answer to subproblem nums[:i+1], where we already used square.

Then to evaluate:

  1. dp(i, 1) we need to look at dp(i - 1, 1) + nums[i] and nums[i] options, because we can not use square.
  2. dp(i, 0) we need to look at 3 options: dp(i-1, 0) + nums[i], where we say that we used square before; dp(i-1, 1) + nums[i]**2, where we say that we use square right now; nums[i]**2, where we take only last number. Note, that there is also option nums[i], but in this case we did not use square at all, so we do not need to take it (moreover nums[i] <= nums[i]**2, so in fact we will always gain if we use square)

Complexity

Time complexity is O(n) as well as space complexity.

Code

class Solution:
    def maxSumAfterOperation(self, nums):
        @lru_cache(None)
        def dp(i, j):
            if i == -1: return 0
            if j == 1: return max(dp(i - 1, j), 0) + nums[i]  #no square
            if j == 0: return max(dp(i-1, j) + nums[i], dp(i-1, j+1) + nums[i]**2, nums[i]**2)

        return max(dp(i, 0) for i in range(len(nums)))