Problem Description

You are given an integer array nums of size n. For each index i in the array (where 0 <= i < n), you need to define a subarray that ends at index i.

The subarray for each index i is defined as nums[start ... i] where:

• start = max(0, i - nums[i])

This means:

• The subarray always ends at index i
• The starting position is calculated by going back nums[i] positions from index i, but it cannot go below index 0
• If i - nums[i] is negative, the subarray starts at index 0
• If i - nums[i] is non-negative, the subarray starts at index i - nums[i]

Your task is to calculate the sum of all elements in each of these subarrays and return the total sum.

For example, if nums = [2, 1, 3]:

• For i = 0: start = max(0, 0 - 2) = 0, subarray is nums[0...0] = [2], sum = 2
• For i = 1: start = max(0, 1 - 1) = 0, subarray is nums[0...1] = [2, 1], sum = 3
• For i = 2: start = max(0, 2 - 3) = 0, subarray is nums[0...2] = [2, 1, 3], sum = 6
• Total sum = 2 + 3 + 6 = 11

Intuition

The key insight is that we need to calculate the sum of elements in multiple subarrays, where each subarray ends at a different index i and starts at max(0, i - nums[i]).

A naive approach would be to iterate through each index i, then iterate from the start position to i to calculate the sum for each subarray. This would result in O(n²) time complexity.

However, we can optimize this using prefix sums. The idea is that if we know the cumulative sum up to any point in the array, we can calculate the sum of any subarray in O(1) time using the formula:

• sum(nums[start...i]) = prefixSum[i+1] - prefixSum[start]

Here's how we arrive at this approach:

1. First, we build a prefix sum array where s[i] represents the sum of all elements from index 0 to index i-1
2. For each index i, we need the sum from start to i, which equals s[i+1] - s[start]
3. The starting position is max(0, i - nums[i])

By precomputing all prefix sums once at the beginning, we can then calculate each subarray sum in constant time. This reduces our overall time complexity from O(n²) to O(n).

The use of accumulate with initial=0 in Python elegantly creates the prefix sum array where s[0] = 0 (sum of zero elements), s[1] = nums[0] (sum up to index 0), s[2] = nums[0] + nums[1] (sum up to index 1), and so on. This indexing makes it convenient to calculate subarray sums using simple subtraction.

Learn more about Prefix Sum patterns.

Solution Approach

The solution implements the prefix sum optimization strategy to efficiently calculate the required subarray sums.

Step 1: Build the Prefix Sum Array

s = list(accumulate(nums, initial=0))

This creates a prefix sum array s where:

• s[0] = 0 (sum of zero elements)
• s[1] = nums[0]
• s[2] = nums[0] + nums[1]
• s[i] = nums[0] + nums[1] + ... + nums[i-1]

The initial=0 parameter ensures the array starts with 0, making it convenient for subarray sum calculations.

Step 2: Calculate Subarray Sums

return sum(s[i + 1] - s[max(0, i - x)] for i, x in enumerate(nums))

For each index i with value x = nums[i]:

• Calculate the starting position: start = max(0, i - x)
• Calculate the subarray sum from start to i: s[i + 1] - s[start]
  • s[i + 1] gives the cumulative sum up to and including index i
  • s[start] gives the cumulative sum up to but not including index start
  • Their difference gives the sum of elements from start to i inclusive

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

1. Prefix sum array: s = [0, 2, 3, 6]
2. For i=0, x=2: start = max(0, 0-2) = 0, sum = s[1] - s[0] = 2 - 0 = 2
3. For i=1, x=1: start = max(0, 1-1) = 0, sum = s[2] - s[0] = 3 - 0 = 3
4. For i=2, x=3: start = max(0, 2-3) = 0, sum = s[3] - s[0] = 6 - 0 = 6
5. Total = 2 + 3 + 6 = 11

Time Complexity: O(n) - We iterate through the array twice (once for prefix sum, once for calculating results) Space Complexity: O(n) - For storing the prefix sum array

Example Walkthrough

Let's walk through a detailed example with nums = [3, 1, 2, 4].

Step 1: Build the prefix sum array

• Start with nums = [3, 1, 2, 4]
• Create prefix sum array: s = [0, 3, 4, 6, 10]
  • s[0] = 0 (no elements)
  • s[1] = 3 (sum up to index 0)
  • s[2] = 3 + 1 = 4 (sum up to index 1)
  • s[3] = 3 + 1 + 2 = 6 (sum up to index 2)
  • s[4] = 3 + 1 + 2 + 4 = 10 (sum up to index 3)

Step 2: Calculate subarray sum for each index

For i = 0 (nums[0] = 3):

• Start position: max(0, 0 - 3) = max(0, -3) = 0
• Subarray: nums[0...0] = [3]
• Sum using prefix: s[1] - s[0] = 3 - 0 = 3

For i = 1 (nums[1] = 1):

• Start position: max(0, 1 - 1) = max(0, 0) = 0
• Subarray: nums[0...1] = [3, 1]
• Sum using prefix: s[2] - s[0] = 4 - 0 = 4

For i = 2 (nums[2] = 2):

• Start position: max(0, 2 - 2) = max(0, 0) = 0
• Subarray: nums[0...2] = [3, 1, 2]
• Sum using prefix: s[3] - s[0] = 6 - 0 = 6

For i = 3 (nums[3] = 4):

• Start position: max(0, 3 - 4) = max(0, -1) = 0
• Subarray: nums[0...3] = [3, 1, 2, 4]
• Sum using prefix: s[4] - s[0] = 10 - 0 = 10

Step 3: Calculate total sum Total = 3 + 4 + 6 + 10 = 23

The key insight is that by precomputing the prefix sums, we can calculate any subarray sum in O(1) time using the formula s[i+1] - s[start], avoiding the need to repeatedly sum the same elements.

Solution Implementation

Time and Space Complexity

Time Complexity: O(n)

• The accumulate function iterates through the array once to build the prefix sum array, which takes O(n) time
• The list comprehension iterates through each element in nums once, performing constant-time operations (s[i + 1] - s[max(0, i - x)]) for each element, which takes O(n) time
• The sum function adds up all the computed values, which also takes O(n) time
• Overall time complexity: O(n) + O(n) + O(n) = O(n)

Space Complexity: O(n)

• The prefix sum array s stores n + 1 elements (including the initial 0), requiring O(n) space
• The list comprehension generates a temporary list of n values before summing them, requiring O(n) space
• Overall space complexity: O(n)

Learn more about how to find time and space complexity quickly.

Common Pitfalls

Pitfall 1: Misunderstanding the Subarray Definition

Issue: A common mistake is misinterpreting how the subarray is defined. Developers might think:

• The subarray has a fixed length of nums[i] elements
• The subarray starts at index i and goes forward
• The calculation is i + nums[i] instead of i - nums[i]

Example of Wrong Understanding:

# WRONG: Thinking subarray starts at i and has length nums[i]
def wrong_solution(nums):
    total = 0
    for i, x in enumerate(nums):
        end = min(len(nums), i + x)  # Wrong direction!
        total += sum(nums[i:end])
    return total

Solution: Remember that:

• The subarray always ENDS at index i
• We go BACKWARDS by nums[i] positions to find the start
• Formula: start = max(0, i - nums[i])

Pitfall 2: Off-by-One Errors with Prefix Sum Indexing

Issue: When using prefix sums, it's easy to confuse the indexing, especially when initial=0 is used.

Common Mistakes:

# WRONG: Incorrect prefix sum indexing
prefix_sums = list(accumulate(nums))  # No initial=0
# Now prefix_sums[i] = sum(nums[0:i+1]) instead of sum(nums[0:i])

# Using wrong indices for calculation
subarray_sum = prefix_sums[i] - prefix_sums[start]  # Wrong!

Solution: When using accumulate(nums, initial=0):

• prefix_sums[k] = sum of first k elements
• To get sum from index start to i inclusive: use prefix_sums[i+1] - prefix_sums[start]
• Always verify with a small example to ensure correct indexing

Pitfall 3: Not Handling Edge Cases

Issue: Failing to properly handle cases where:

• nums[i] is 0 (subarray is just the single element at index i)
• nums[i] is greater than i (subarray starts from index 0)
• Array contains negative numbers (if allowed by constraints)

Example of Incomplete Handling:

# WRONG: Might fail for edge cases
def incomplete_solution(nums):
    total = 0
    for i in range(len(nums)):
        start = i - nums[i]  # Forgot max(0, ...)!
        # This could give negative start index
        total += sum(nums[start:i+1])
    return total

Solution: Always use max(0, i - nums[i]) to ensure the start index never goes below 0, and test with edge cases like nums = [0, 5, 1] where the second element would try to go back 5 positions from index 1.