Maximum Subarray Sum After One Operation

medium array dp

Problem

You are given an integer array nums. You must replace exactly one element nums[i] with its square (nums[i] * nums[i]). After performing that single operation, return the maximum possible sum of a contiguous subarray (the subarray must contain at least one element).

Inputnums = [2, -1, -4, -3]

Output17

Square the -4 to get 16, then take the subarray [2, -1, 16] whose sum is 2 - 1 + 16 = 17.

nums (comma-separated, signed)

def max_sum_after_operation(nums):
    no_sq = 0
    used = 0
    best = nums[0] * nums[0]
    for x in nums:
        used = max(no_sq + x * x, used + x, x * x)
        no_sq = max(no_sq + x, x)
        best = max(best, used)
    return best

function maxSumAfterOperation(nums) {
  let noSq = 0;
  let used = 0;
  let best = nums[0] * nums[0];
  for (const x of nums) {
    used = Math.max(noSq + x * x, used + x, x * x);
    noSq = Math.max(noSq + x, x);
    best = Math.max(best, used);
  }
  return best;
}

class Solution {
    public int maxSumAfterOperation(int[] nums) {
        int noSq = 0;
        int used = 0;
        int best = nums[0] * nums[0];
        for (int x : nums) {
            used = Math.max(Math.max(noSq + x * x, used + x), x * x);
            noSq = Math.max(noSq + x, x);
            best = Math.max(best, used);
        }
        return best;
    }
}

int maxSumAfterOperation(vector<int>& nums) {
    int noSq = 0;
    int used = 0;
    int best = nums[0] * nums[0];
    for (int x : nums) {
        used = max({ noSq + x * x, used + x, x * x });
        noSq = max(noSq + x, x);
        best = max(best, used);
    }
    return best;
}

Explanation

This is Kadane's algorithm with a twist: we are allowed (and required) to square one element somewhere inside the subarray. The natural way to handle a one-time choice is to carry two running sums as we scan left to right.

Let noSq be the best subarray sum ending at the current position where we have not used the square yet, and let used be the best subarray sum ending here where we have already squared one element.

For each value x, the plain Kadane update gives noSq = max(noSq + x, x) — either extend the previous run or start fresh at x. For the squared state we have three choices: extend a not-yet-squared run and square this element (noSq + x*x), extend an already-squared run with the plain value (used + x), or start a brand-new run that begins by squaring this element (x*x). We take the largest of the three.

Because the operation is mandatory, the answer must come from the used state, so we track best as the maximum used ever seen.

Example: nums = [2, -1, -4, -3]. Scanning along, the squared state peaks at the -4 step: extending the run [2, -1] (sum 1) and squaring -4 into 16 gives 1 + 16 = 17. No later step beats it, so the answer is 17.

Time: O(n) Space: O(1)