题目链接:https://leetcode.com/problems/sum-of-all-odd-length-subarrays/
要求:计算所有奇数长度的子数组的和。
提到子数组,自然先想到前缀和。
尝试用前缀和来解答。
数组记做A,长度为n,假设子数组长度为i(i=1、3、5、…),则每个子数组的和就是
A[0] + A[1] + … + A[i-1]
A[1] + A[2] + … + A[i]
…
A[n-i] + A[n-i+1] + … + A[n-1]
前缀和数组记做S,上述算式可以演化为
S[i-1]
S[i] - S[0]
…
S[n-1] - S[n-i-1]
那么最终结果就是
(S[i-1] + S[i] + … + S[n-1]) - (S[0] + S[1] + … + S[n-i-1])
这里一看,又是一个前缀和。
假设数组S的前缀和数组记做P,则结果就是
P[n-1] - P[i-2] - P[n-i-1]
最终代码如下:
func sumOddLengthSubarrays(arr []int) int {
lth := len(arr)
if lth == 1 {
return arr[0]
}
prefixSum := runningSum(arr)
prefixSum2 := runningSum(prefixSum)
res := prefixSum[lth-1]
for i := 3; i <= lth-1; i += 2 {
res += prefixSum2[lth-1] - prefixSum2[i-2] - prefixSum2[lth-i-1]
}
if lth % 2 == 1 {
res += prefixSum[lth-1]
}
return res
}
func runningSum(nums []int) []int {
lth := len(nums)
res := make([]int, lth)
res[0] = nums[0]
for i := 1; i < lth; i++ {
res[i] = res[i-1] + nums[i]
}
return res
}