Program to find out the sum of the maximum subarray after a operation in Python

Given an array of integers, we can perform one operation where we replace any element array[i] with its squared value array[i] * array[i]. Our goal is to find the maximum possible subarray sum after performing this operation.

So, if the input is like array = [4, 1, -2, -1], then the output will be 17.

If we replace the value in array[0] with its squared value, the array becomes [16, 1, -2, -1]. The maximum subarray from this is [16, 1], which has the sum 16 + 1 = 17.

Algorithm

We use dynamic programming with two arrays:

• dp1 − stores maximum subarray sum without using the square operation
• dp2 − stores maximum subarray sum after using the square operation at most once

For each element, we consider three possibilities:

• Continue previous subarray without squaring current element
• Start new subarray by squaring current element
• Square current element and add to previous subarray (if operation not used)

Implementation

def solve(array):
    # dp1: max subarray sum without using square operation
    dp1 = [float('-inf')]
    # dp2: max subarray sum using square operation at most once
    dp2 = [float('-inf')]
    
    for num in array:
        # Without squaring: continue previous or start new
        dp1.append(max(dp1[-1] + num, num))
        
        # With squaring: three options
        # 1. Square current and add to subarray before squaring
        # 2. Start new subarray with squared current element
        # 3. Continue previous subarray (already used square operation)
        dp2.append(max(dp1[-2] + num**2, num**2, dp2[-1] + num))
    
    return max(dp2)

# Test the function
result = solve([4, 1, -2, -1])
print(f"Maximum subarray sum: {result}")

Maximum subarray sum: 17

How It Works

Let's trace through the example [4, 1, -2, -1]:

def solve_with_trace(array):
    dp1 = [float('-inf')]
    dp2 = [float('-inf')]
    
    print("Step by step execution:")
    print("dp1 (no square):", dp1)
    print("dp2 (with square):", dp2)
    print()
    
    for i, num in enumerate(array):
        # Update dp1
        dp1.append(max(dp1[-1] + num, num))
        
        # Update dp2
        option1 = dp1[-2] + num**2  # Square current, add to previous
        option2 = num**2            # Start new with squared
        option3 = dp2[-1] + num     # Continue with current
        
        dp2.append(max(option1, option2, option3))
        
        print(f"After processing {num}:")
        print(f"  dp1: {dp1}")
        print(f"  dp2: {dp2}")
        print(f"  Options for dp2: {option1}, {option2}, {option3}")
        print()
    
    return max(dp2)

solve_with_trace([4, 1, -2, -1])

Step by step execution:
dp1 (no square): [-inf]
dp2 (with square): [-inf]

After processing 4:
  dp1: [-inf, 4]
  dp2: [-inf, 16]
  Options for dp2: 16, 16, -inf

After processing 1:
  dp1: [-inf, 4, 5]
  dp2: [-inf, 16, 17]
  Options for dp2: 5, 1, 17

After processing -2:
  dp1: [-inf, 4, 5, 3]
  dp2: [-inf, 16, 17, 19]
  Options for dp2: 9, 4, 15

After processing -1:
  dp1: [-inf, 4, 5, 3, 2]
  dp2: [-inf, 16, 17, 19, 18]
  Options for dp2: 4, 1, 18

Key Points

• We can use the square operation at most once in the entire array
• The subarray must be non-empty
• Dynamic programming tracks both possibilities: with and without using the square operation
• Time complexity: O(n), Space complexity: O(n)

Conclusion

This dynamic programming approach efficiently finds the maximum subarray sum after performing at most one square operation. The key insight is tracking two states: maximum sum without squaring and maximum sum with squaring used at most once.