Problem

You are given an integer array nums. You must perform exactly one operation where you can replace one element nums[i] with nums[i] * nums[i].

Return themaximum possible subarray sum after exactly one operation. The subarray must be non-empty.

Examples

Example 1:

1
2
3
Input: nums = [2,-1,-4,-3]
Output: 17
Explanation: You can perform the operation on index 2 (0-indexed) to make nums = [2,-1,**16** ,-3]. Now, the maximum subarray sum is 2 + -1 + 16 = 17.

Example 2:

1
2
3
Input: nums = [1,-1,1,1,-1,-1,1]
Output: 4
Explanation: You can perform the operation on index 1 (0-indexed) to make nums = [1,**1** ,1,1,-1,-1,1]. Now, the maximum subarray sum is 1 + 1 + 1 + 1 = 4.

Constraints:

  • 1 <= nums.length <= 10^5
  • -10^4 <= nums[i] <= 10^4

Solution

Method 1 – Dynamic Programming (Kadane’s Variant)

Intuition

We want the maximum subarray sum after squaring exactly one element. For each position, we can either have not used the operation yet, or have already used it. We use dynamic programming to track both cases as we scan the array.

Approach

  1. Let no_op be the maximum subarray sum ending at the current index without using the operation.
  2. Let with_op be the maximum subarray sum ending at the current index having used the operation.
  3. For each index:
    • Update no_op as the max of nums[i] and no_op + nums[i] (Kadane’s algorithm).
    • Update with_op as the max of:
      • no_op + nums[i]*nums[i] (use operation at this index)
      • with_op + nums[i] (already used operation before)
      • nums[i]*nums[i] (start new subarray with operation at this index)
  4. Track the maximum value of with_op across all indices.

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
class Solution {
public:
    int maxSumAfterOperation(vector<int>& nums) {
        int n = nums.size();
        int no_op = 0, with_op = 0, ans = INT_MIN;
        for (int i = 0; i < n; ++i) {
            with_op = max({no_op + nums[i]*nums[i], with_op + nums[i], nums[i]*nums[i]});
            no_op = max(no_op + nums[i], nums[i]);
            ans = max(ans, with_op);
        }
        return ans;
    }
};
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
func maxSumAfterOperation(nums []int) int {
    noOp, withOp := 0, 0
    ans := nums[0]*nums[0]
    for _, x := range nums {
        withOp = max(noOp + x*x, withOp + x, x*x)
        noOp = max(noOp + x, x)
        ans = max(ans, withOp)
    }
    return ans
}
func max(a, b int) int { if a > b { return a }; return b }
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
class Solution {
    public int maxSumAfterOperation(int[] nums) {
        int noOp = 0, withOp = 0, ans = Integer.MIN_VALUE;
        for (int x : nums) {
            withOp = Math.max(Math.max(noOp + x*x, withOp + x), x*x);
            noOp = Math.max(noOp + x, x);
            ans = Math.max(ans, withOp);
        }
        return ans;
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
class Solution {
    fun maxSumAfterOperation(nums: IntArray): Int {
        var noOp = 0; var withOp = 0; var ans = Int.MIN_VALUE
        for (x in nums) {
            withOp = maxOf(noOp + x*x, withOp + x, x*x)
            noOp = maxOf(noOp + x, x)
            ans = maxOf(ans, withOp)
        }
        return ans
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
class Solution:
    def maxSumAfterOperation(self, nums: list[int]) -> int:
        no_op = 0
        with_op = 0
        ans = float('-inf')
        for x in nums:
            with_op = max(no_op + x*x, with_op + x, x*x)
            no_op = max(no_op + x, x)
            ans = max(ans, with_op)
        return ans
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
impl Solution {
    pub fn max_sum_after_operation(nums: Vec<i32>) -> i32 {
        let (mut no_op, mut with_op, mut ans) = (0, 0, i32::MIN);
        for &x in &nums {
            with_op = (no_op + x*x).max(with_op + x).max(x*x);
            no_op = (no_op + x).max(x);
            ans = ans.max(with_op);
        }
        ans
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
class Solution {
    maxSumAfterOperation(nums: number[]): number {
        let noOp = 0, withOp = 0, ans = -Infinity;
        for (const x of nums) {
            withOp = Math.max(noOp + x*x, withOp + x, x*x);
            noOp = Math.max(noOp + x, x);
            ans = Math.max(ans, withOp);
        }
        return ans;
    }
}

Complexity

  • ⏰ Time complexity: O(n), where n is the length of nums. We scan the array once.
  • 🧺 Space complexity: O(1), only a few variables are used.