LeetCode 3427 - Sum of Variable Length Subarrays

LeetCode Problem 3427

Difficulty: 🟢 Easy
Topics: Array, Prefix Sum

Solution

Problem Understanding

The problem provides an integer array nums of length n. For each index i, we are asked to consider a subarray that ends at i but starts at start = max(0, i - nums[i]). This effectively means the subarray has a variable length depending on the value of nums[i] but never extends beyond the beginning of the array. The task is to compute the sum of all these subarrays for every index and return the total sum.

The input nums is constrained such that 1 <= n <= 100 and 1 <= nums[i] <= 1000. These constraints imply that the array is relatively small, so even an O(n²) solution would be acceptable. Important edge cases include small arrays (length 1), large values of nums[i] that cover the entire array, and cases where nums[i] is equal to the index (so the subarray starts at index 0).

Approaches

A brute-force approach would iterate through each index i, determine the subarray starting at start = max(0, i - nums[i]) and ending at i, and sum all the elements in that subarray. This works correctly but requires iterating over potentially large subarrays for each index, resulting in O(n²) time complexity. Given the constraints, it is acceptable but not elegant.

A more optimal approach leverages a prefix sum array to compute subarray sums in O(1) time. A prefix sum array prefix is constructed such that prefix[i] stores the sum of elements from nums[0] to nums[i]. Then the sum of any subarray nums[start ... i] is simply prefix[i] - prefix[start - 1] if start > 0, or prefix[i] if start == 0. This reduces the computation of each subarray sum to O(1), and the total algorithm runs in O(n) time.

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(n²)	O(1)	Directly sum each subarray for every index
Optimal	O(n)	O(n)	Use prefix sum array to compute each subarray sum efficiently

Algorithm Walkthrough

1. Initialize a prefix sum array prefix of length n.
2. Compute the prefix sum for the first element: prefix[0] = nums[0].
3. For each index i from 1 to n-1, compute prefix[i] = prefix[i-1] + nums[i].
4. Initialize a variable total_sum to 0.
5. Iterate through each index i from 0 to n-1.
6. Compute the starting index of the subarray as start = max(0, i - nums[i]).
7. Calculate the subarray sum: if start > 0, then sub_sum = prefix[i] - prefix[start - 1]; otherwise, sub_sum = prefix[i].
8. Add sub_sum to total_sum.
9. Return total_sum.

Why it works: The prefix sum ensures that for any subarray nums[start ... i], we can compute its sum in O(1). By iterating through all indices, we guarantee that each subarray sum is included exactly once in the total sum.

Python Solution

from typing import List

class Solution:
    def subarraySum(self, nums: List[int]) -> int:
        n = len(nums)
        prefix = [0] * n
        prefix[0] = nums[0]
        for i in range(1, n):
            prefix[i] = prefix[i-1] + nums[i]
        
        total_sum = 0
        for i in range(n):
            start = max(0, i - nums[i])
            if start > 0:
                total_sum += prefix[i] - prefix[start - 1]
            else:
                total_sum += prefix[i]
        
        return total_sum

This implementation first constructs the prefix sum array, which allows each subarray sum to be computed in constant time. Then, for each index, it determines the subarray start index and computes the subarray sum using the prefix sum array, accumulating the total.

Go Solution

func subarraySum(nums []int) int {
    n := len(nums)
    prefix := make([]int, n)
    prefix[0] = nums[0]
    for i := 1; i < n; i++ {
        prefix[i] = prefix[i-1] + nums[i]
    }

    totalSum := 0
    for i := 0; i < n; i++ {
        start := i - nums[i]
        if start < 0 {
            start = 0
        }
        if start > 0 {
            totalSum += prefix[i] - prefix[start-1]
        } else {
            totalSum += prefix[i]
        }
    }

    return totalSum
}

In Go, we use slices for the prefix array. Edge cases such as start < 0 are handled by resetting start to 0. The logic closely mirrors the Python version.

Worked Examples

Example 1: nums = [2,3,1]

i	nums[i]	start	Subarray	Subarray Sum	Total Sum
0	2	0	[2]	2	2
1	3	0	[2,3]	5	7
2	1	1	[3,1]	4	11

Example 2: nums = [3,1,1,2]

i	nums[i]	start	Subarray	Subarray Sum	Total Sum
0	3	0	[3]	3	3
1	1	0	[3,1]	4	7
2	1	1	[1,1]	2	9
3	2	1	[1,1,2]	4	13

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n)	Prefix sum computation takes O(n), and iterating to compute each subarray sum also takes O(n)
Space	O(n)	Storing the prefix sum array of length n

The prefix sum array ensures each subarray sum is computed in constant time, resulting in linear total computation.

Test Cases

# Provided examples
assert Solution().subarraySum([2,3,1]) == 11  # example 1
assert Solution().subarraySum([3,1,1,2]) == 13  # example 2

# Edge cases
assert Solution().subarraySum([1]) == 1  # single element
assert Solution().subarraySum([1000]*100) == sum([sum([1000]*min(i+1,1000)) for i in range(100)])  # large values

# Increasing sequence
assert Solution().subarraySum([1,2,3,4,5]) == 35

# Decreasing sequence
assert Solution().subarraySum([5,4,3,2,1]) == 35

Test	Why
[2,3,1]	Validates correct subarray computation and sum aggregation
[3,1,1,2]	Tests variable-length subarrays
[1]	Minimal input edge case
[1000]*100	Large identical values to stress summation correctness
[1,2,3,4,5]	Increasing sequence for general correctness
[5,4,3,2,1]	Decreasing sequence for general correctness

Edge Cases

One important edge case is a single-element array, where the subarray sum is trivially the element itself. Another is when nums[i] is larger than the index i, causing start to be negative; we handle this with max(0, i - nums[i]) to avoid index errors. Finally, arrays with large repeated numbers can cause integer overflow in some languages, so the solution must use a datatype capable of handling large sums, which Python does automatically with its integers. In Go, the standard int type is sufficient for the given constraints.