小而美的算法技巧：前缀和数组

小而美的算法技巧：前缀和数组

前缀和主要适用的场景是原始数组不会被修改的情况下，频繁查询某个区间的累加和。

一维数组中的前缀和

https://leetcode-cn.com/problems/range-sum-query-immutable/

前缀和技巧：提前计算数组从 0 开始的累加和并储存，在需要获取一个区间内的和时，直接在累加和数组中相减即可得到，不需要每次去重新遍历数组。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
type NumArray struct {
    prefixSum []int
}


func Constructor(nums []int) NumArray {
    a := new(NumArray)
    prefixSum := make([]int, len(nums) + 1)
    // 计算前缀和
    for i := 1; i <= len(nums); i++ {
        prefixSum[i] = prefixSum[i - 1] + nums[i - 1]
    }
    a.prefixSum = prefixSum
    return *a
}


func (this *NumArray) SumRange(left int, right int) int {
    return this.prefixSum[right + 1] - this.prefixSum[left]
}

二维矩阵中的前缀和

https://leetcode-cn.com/problems/range-sum-query-2d-immutable/

类似于上一题，不过是二维的。对于二维情况稍微复杂一点，但原理还是一样的：通过几个子矩阵前缀和的加减运算得出任意区域矩阵的和。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
type NumMatrix struct {
    prefixMatrix [][]int
}


func Constructor(matrix [][]int) NumMatrix {
    m, n := len(matrix), len(matrix[0])
    prefix := make([][]int, m + 1)
    for i := range prefix {
        prefix[i] = make([]int, n + 1)
    }

    for i := 1; i <= m; i++ {
        for j := 1; j <= n; j++ {
            prefix[i][j] = prefix[i-1][j] + prefix[i][j-1] - prefix[i-1][j-1] + matrix[i-1][j-1]
        }
    }
    return NumMatrix{prefix}
}


func (this *NumMatrix) SumRegion(row1 int, col1 int, row2 int, col2 int) int {
    return this.prefixMatrix[row2+1][col2+1] - this.prefixMatrix[row1][col2+1] - this.prefixMatrix[row2+1][col1] + this.prefixMatrix[row1][col1]
}

和为 k 的子数组

类似于两数之和，计算出前缀数组后不用再去遍历一遍，只需要寻找它减去 k 以后的前缀和是否已经在哈希表中。这样的前缀和有几个，就有几个能凑出和为 k 的子数组，直接在结果加上即可。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
func subarraySum(nums []int, k int) int {
    m := make(map[int]int)
    m[0] = 1  // base case
    
    res, sum:= 0, 0
    for i := 0; i < len(nums); i++ {
        sum += nums[i]
        start := sum - k
        if val, ok := m[start]; ok {
            res += val
        }
        m[sum] += 1
    }

    return res
}