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.
We use prefix sums over nums to quickly compute the sum of each subarray. For each index i (where 0 <= i < n), compute start = max(0, i - nums[i]) and add the subarray sum prefix[i+1] - prefix[start] to the running total ans.
classSolution {
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;
}
};
1
2
3
4
5
6
7
8
9
10
11
12
13
classSolution {
publicintsumOfVariableLengthSubarrays(int[] nums) {
int n = nums.length;
int[] prefix =newint[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;
}
}
1
2
3
4
5
6
7
8
9
10
11
classSolution:
defsumOfVariableLengthSubarrays(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 =0for i in range(n):
start = max(0, i - nums[i])
ans += prefix[i+1] - prefix[start]
return ans
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
impl Solution {
pubfnsum_of_variable_length_subarrays(nums: Vec<i32>) -> i32 {
let n = nums.len();
letmut prefix =vec![0; n+1];
for i in0..n {
prefix[i+1] = prefix[i] + nums[i];
}
letmut ans =0;
for i in0..n {
let start =if i >= nums[i] asusize { i - nums[i] asusize } else { 0 };
ans += prefix[i+1] - prefix[start];
}
ans
}
}