2025-07-30:变长子数组求和。用go语言,给定一个长度为 n 的整数数组 nums。对于数组中的每个位置 i(范围是 0 到 n-1),我们定义一个子数组,区间为 nums[start ... i],其中 start 等于 max(0, i - nums[i])。
任务是计算并返回对于每个位置 i,所对应的子数组内所有元素的累加和。
简而言之,就是对于每个元素,根据它的值确定子数组的起始位置,然后求该子数组的元素和。
1 <= n == nums.length <= 100。
1 <= nums[i] <= 1000。
输入:nums = [2,3,1]。
输出:11。
解释:
下标 i | 子数组 | 和 |
---|---|---|
0 | nums[0] = [2] | 2 |
1 | nums[0 ... 1] = [2, 3] | 5 |
2 | nums[1 ... 2] = [3, 1] | 4 |
总和 | 11 |
总和为 11 。因此,输出 11 。
题目来自力扣3427。
diff
len(nums) + 1
的整数数组 diff
,初始时全部为0。对 nums
中每个元素 nums[i]
:
start = max(0, i - nums[i])
diff[start]++
:表示从位置 start
开始,有一个新的子数组加进来diff[i+1]--
:表示从位置 i+1
开始,前面的那个子数组不再计入(因为子数组截止到 i)换句话说,用差分数组标记了每个子数组所覆盖的区间。
定义变量 sd = 0
,它表示当前索引 i 处子数组的数量。
遍历数组索引 i:
sd += diff[i]
:通过累加差分,得到当前位置 i 被多少个子数组覆盖(“活跃”的子数组个数)nums[i]
,它出现于 sd
个子数组里。nums[i]
的总加和值贡献就是 nums[i] * sd
。将这些贡献累加到变量 ans
中。
ans
这种做法避免了重复计算每个子数组内元素,从而提高效率。
n+1
的差分数组 diff
总结:
该算法利用差分数组技巧,从而避免了嵌套循环的暴力求和(O(n^2)),效率较高。
.
package main
import (
"fmt"
)
func subarraySum(nums []int) (ans int) {
diff := make([]int, len(nums)+1)
for i, num := range nums {
diff[max(i-num, 0)]++
diff[i+1]--
}
sd := 0
for i, x := range nums {
sd += diff[i]
ans += x * sd
}
return
}
func main() {
nums := []int{2, 3, 1}
result := subarraySum(nums)
fmt.Println(result)
}
.
# -*-coding:utf-8-*-
def subarray_sum(nums):
n = len(nums)
diff = [0] * (n + 1)
for i, num in enumerate(nums):
start = max(i - num, 0)
diff[start] += 1
diff[i + 1] -= 1
ans = 0
sd = 0
for i, x in enumerate(nums):
sd += diff[i]
ans += x * sd
return ans
if __name__ == "__main__":
nums = [2, 3, 1]
result = subarray_sum(nums)
print(result)