Problem

You are given an integer array nums of size n. For each index i where 0 <= i < n, define a subarray nums[start ... i] where start = max(0, i - nums[i]).

Return the total sum of all elements from the subarray defined for each index in the array.

Example 1

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

Input: nums = [2,3,1]
Output: 11
Explanation:
i | Subarray | Sum  
---|---|---  
0 | `nums[0] = [2]` | 2  
1 | `nums[0 ... 1] = [2, 3]` | 5  
2 | `nums[1 ... 2] = [3, 1]` | 4  
**Total Sum** |   | 11  
The total sum is 11. Hence, 11 is the output.

Example 2

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

Input: nums = [3,1,1,2]
Output: 13
Explanation:
i | Subarray | Sum  
---|---|---  
0 | `nums[0] = [3]` | 3  
1 | `nums[0 ... 1] = [3, 1]` | 4  
2 | `nums[1 ... 2] = [1, 1]` | 2  
3 | `nums[1 ... 3] = [1, 1, 2]` | 4  
**Total Sum** |   | 13  
The total sum is 13. Hence, 13 is the output.

Constraints

  • 1 <= n == nums.length <= 100
  • 1 <= nums[i] <= 1000

Examples

Solution

Approach

We use prefix sums to efficiently compute the sum of each subarray. For each index i, the subarray is from start = max(0, i - nums[i]) to i. The sum is prefix[i+1] - prefix[start].

Code

C++
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

class Solution {
public:
    int sumOfVariableLengthSubarrays(vector<int>& nums) {
        int n = nums.size();
        vector<int> prefix(n+1, 0);
        for (int i = 0; i < n; ++i) prefix[i+1] = prefix[i] + nums[i];
        int ans = 0;
        for (int i = 0; i < n; ++i) {
            int start = max(0, i - nums[i]);
            ans += prefix[i+1] - prefix[start];
        }
        return ans;
    }
};
Java
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

class Solution {
    public int sumOfVariableLengthSubarrays(int[] nums) {
        int n = nums.length;
        int[] prefix = new int[n+1];
        for (int i = 0; i < n; ++i) prefix[i+1] = prefix[i] + nums[i];
        int ans = 0;
        for (int i = 0; i < n; ++i) {
            int start = Math.max(0, i - nums[i]);
            ans += prefix[i+1] - prefix[start];
        }
        return ans;
    }
}
Python
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

class Solution:
    def sumOfVariableLengthSubarrays(self, nums: list[int]) -> int:
        n = len(nums)
        prefix = [0] * (n+1)
        for i in range(n):
            prefix[i+1] = prefix[i] + nums[i]
        ans = 0
        for i in range(n):
            start = max(0, i - nums[i])
            ans += prefix[i+1] - prefix[start]
        return ans
Rust
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15

impl Solution {
    pub fn sum_of_variable_length_subarrays(nums: Vec<i32>) -> i32 {
        let n = nums.len();
        let mut prefix = vec![0; n+1];
        for i in 0..n {
            prefix[i+1] = prefix[i] + nums[i];
        }
        let mut ans = 0;
        for i in 0..n {
            let start = if i >= nums[i] as usize { i - nums[i] as usize } else { 0 };
            ans += prefix[i+1] - prefix[start];
        }
        ans
    }
}
TypeScript
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

function sumOfVariableLengthSubarrays(nums: number[]): number {
    const n = nums.length;
    const prefix = new Array(n+1).fill(0);
    for (let i = 0; i < n; ++i) prefix[i+1] = prefix[i] + nums[i];
    let ans = 0;
    for (let i = 0; i < n; ++i) {
        const start = Math.max(0, i - nums[i]);
        ans += prefix[i+1] - prefix[start];
    }
    return ans;
}

Explanation

We use prefix sums to quickly compute the sum of each subarray. For each index, we add the sum of its subarray to the answer.

Complexity

  • ⏰ Time complexity: O(n)
  • 🧺 Space complexity: O(n)