Taro Logo

Maximum Sum of 3 Non-Overlapping Subarrays

Hard
Google logo
Google
Topics:
ArraysSliding WindowsDynamic Programming

Given an integer array nums and an integer k, find three non-overlapping subarrays of length k with maximum sum and return them.

Return the result as a list of indices representing the starting position of each interval (0-indexed). If there are multiple answers, return the lexicographically smallest one.

For example:

  • nums = [1,2,1,2,6,7,5,1], k = 2

    Output: [0,3,5]

    Explanation: Subarrays [1, 2], [2, 6], [7, 5] correspond to the starting indices [0, 3, 5]. We could have also taken [2, 1], but an answer of [1, 3, 5] would be lexicographically larger.

  • nums = [1,2,1,2,1,2,1,2,1], k = 2

    Output: [0,2,4]

Explain your solution. What is the Big(O) run-time and space usage of your solution? What are some edge cases?

Solution


Brute Force Solution

A naive approach would be to iterate through all possible combinations of three non-overlapping subarrays of length k and compute the sum of each combination. We keep track of the maximum sum encountered so far and the corresponding indices.

Code:

def max_sum_three_subarrays_brute_force(nums, k):
    n = len(nums)
    max_sum = 0
    result = []

    for i in range(n - 3 * k + 1):
        for j in range(i + k, n - 2 * k + 1):
            for l in range(j + k, n - k + 1):
                sum1 = sum(nums[i:i+k])
                sum2 = sum(nums[j:j+k])
                sum3 = sum(nums[l:l+k])
                current_sum = sum1 + sum2 + sum3

                if current_sum > max_sum:
                    max_sum = current_sum
                    result = [i, j, l]

    return result

Time Complexity: O(n^3), where n is the length of the input array. This is due to the three nested loops. Space Complexity: O(1), as we are only using a constant amount of extra space.

Optimal Solution: Sliding Window and Dynamic Programming

A more efficient approach involves using a sliding window to compute the sums of all subarrays of length k. Then, we can use dynamic programming-like ideas to find the optimal combination of three subarrays.

  1. Calculate subarray sums: First, calculate the sum of all subarrays of length k and store them in an array sums.
  2. Left array: Create a left array where left[i] stores the index of the maximum sum subarray of length k from 0 to i. This can be done in one pass.
  3. Right array: Similarly, create a right array where right[i] stores the index of the maximum sum subarray of length k from i to n - k -1. This can be done in one pass.
  4. Iterate middle subarray: Iterate through the possible starting indices of the middle subarray. For each middle subarray starting at index j, combine it with the best left subarray (from left[j-k]) and the best right subarray (from right[j+k]). Keep track of the best combination seen so far.

Code:

def max_sum_three_subarrays(nums, k):
    n = len(nums)
    sums = []
    window_sum = sum(nums[:k])
    sums.append(window_sum)

    for i in range(k, n):
        window_sum += nums[i] - nums[i - k]
        sums.append(window_sum)

    left = [0] * (n - k + 1)
    best_index = 0
    for i in range(1, n - k + 1):
        if sums[i] > sums[best_index]:
            best_index = i
        left[i] = best_index

    right = [n - k] * (n - k + 1)
    best_index = n - k
    for i in range(n - k - 1, -1, -1):
        if sums[i] >= sums[best_index]:
            best_index = i
        right[i] = best_index

    max_sum = 0
    result = []

    for j in range(k, n - 2 * k + 1):
        left_index = left[j - k]
        right_index = right[j + k]
        current_sum = sums[left_index] + sums[j] + sums[right_index]

        if current_sum > max_sum:
            max_sum = current_sum
            result = [left_index, j, right_index]

    return result

Time Complexity: O(n), where n is the length of the input array. We iterate through the array a few times to compute prefix sums and the left and right arrays. Space Complexity: O(n), as we are using extra space to store the sums, left, and right arrays.

Edge Cases

  • Empty array: Handle the case where the input array is empty. Although the constraints state 1 <= nums.length <= 2 * 10^4, it's good practice to consider such cases.
  • Invalid k: Handle the case where k is zero or larger than floor(nums.length / 3). Again, the problem statement has constraints for this, but error checking is good.
  • Multiple answers: The problem specifies returning the lexicographically smallest answer. The provided solution inherently returns the lexicographically smallest result because it iterates from left to right and updates the result only when a strictly larger sum is found. In the right array the index i is only updated if it is greater or equal than the best index so far, thus guaranteeing the lexicographical order.