小而美的算法技巧：前缀和数组
Info
已完成网站教程、网站习题、配套插件中所有多语言代码的校准，解决了之前 chatGPT 翻译可能出错的问题~
读完本文，你不仅学会了算法套路，还可以顺便解决如下题目：
| LeetCode | Difficulty |
|---|---|
| 303. Range Sum Query - Immutable | 🟢 |
| 304. Range Sum Query 2D - Immutable | 🟠 |
Tip
本文有视频版：Prefix Sum/Difference Array Techniques Explained。建议关注我的 B 站账号，我会用视频领读的方式带大家学习那些稍有难度的算法技巧。
The prefix sum technique is useful for quickly and frequently calculating the sum of elements in an index range.
Prefix Sum in One-Dimensional Arrays
Let's start with an example, LeetCode problem 303 "Range Sum Query - Immutable", which asks you to calculate the sum of elements in a subarray. This is a classic prefix sum problem:
303. Range Sum Query - Immutable | 力扣 | LeetCode |
Given an integer array nums, handle multiple queries of the following type:
- Calculate the sum of the elements of
numsbetween indicesleftandrightinclusive whereleft <= right.
Implement the NumArray class:
NumArray(int[] nums)Initializes the object with the integer arraynums.int sumRange(int left, int right)Returns the sum of the elements ofnumsbetween indicesleftandrightinclusive (i.e.nums[left] + nums[left + 1] + ... + nums[right]).
Example 1:
Input ["NumArray", "sumRange", "sumRange", "sumRange"] [[[-2, 0, 3, -5, 2, -1]], [0, 2], [2, 5], [0, 5]] Output [null, 1, -1, -3] Explanation NumArray numArray = new NumArray([-2, 0, 3, -5, 2, -1]); numArray.sumRange(0, 2); // return (-2) + 0 + 3 = 1 numArray.sumRange(2, 5); // return 3 + (-5) + 2 + (-1) = -1 numArray.sumRange(0, 5); // return (-2) + 0 + 3 + (-5) + 2 + (-1) = -3
Constraints:
1 <= nums.length <= 104-105 <= nums[i] <= 1050 <= left <= right < nums.length- At most
104calls will be made tosumRange.
// The problem requires you to implement such a class
class NumArray {
public NumArray(int[] nums) {}
// Query the cumulative sum of the closed interval [left, right]
public int sumRange(int left, int right) {}
}// The problem requires you to implement such a class
class NumArray {
public:
NumArray(vector<int>& nums) {}
// Query the accumulated sum of the closed interval [left, right]
int sumRange(int left, int right) {}
};# The problem requires you to implement such a class
class NumArray:
def __init__(self, nums: List[int]):
pass
# Query the accumulated sum of the closed interval [left, right]
def sumRange(self, left: int, right: int) -> int:
pass// The problem requires you to implement such a class
type NumArray struct {}
func Constructor(nums []int) NumArray {}
// Query the cumulative sum of the closed interval [left, right]
func (this *NumArray) SumRange(left int, right int) int {}// The problem requires you to implement such a class
var NumArray = function(nums) {
this.nums = nums;
// Query the accumulated sum of the closed interval [left, right]
this.sumRange = function(left, right) {}
};The sumRange function needs to calculate and return the sum of elements within a given index range. People who haven't learned about prefix sums might write code like this:
class NumArray {
private int[] nums;
public NumArray(int[] nums) {
this.nums = nums;
}
public int sumRange(int left, int right) {
int res = 0;
for (int i = left; i <= right; i++) {
res += nums[i];
}
return res;
}
}class NumArray {
private:
vector<int> nums;
public:
NumArray(vector<int>& nums) {
this->nums = nums;
}
int sumRange(int left, int right) {
int res = 0;
for (int i = left; i <= right; i++) {
res += nums[i];
}
return res;
}
};class NumArray:
def __init__(self, nums: List[int]):
self.nums = nums
def sumRange(self, left: int, right: int) -> int:
res = 0
for i in range(left, right+1):
res += self.nums[i]
return restype NumArray struct {
nums []int
}
func Constructor(nums []int) NumArray {
return NumArray{nums: nums}
}
func (this *NumArray) SumRange(left int, right int) int {
res := 0
for i := left; i <= right; i++{
res += this.nums[i]
}
return res
}var NumArray = function(nums) {
this.nums = nums;
}
NumArray.prototype.sumRange = function(left, right) {
let res = 0;
for (let i = left; i <= right; i++) {
res += this.nums[i];
}
return res;
}This approach can achieve the desired result, but it is very inefficient because the sumRange method will be called frequently, and its worst-case time complexity is O(N), where N represents the length of the nums array.
The optimal solution for this problem is to use the prefix sum technique, which reduces the time complexity of the sumRange function to O(1). In simple terms, this means avoiding the use of a for loop inside sumRange. How do we do this?
Let's directly look at the code implementation:
class NumArray {
// prefix sum array
private int[] preSum;
// input an array to construct the prefix sum
public NumArray(int[] nums) {
// preSum[0] = 0, convenient for calculating the cumulative sum
preSum = new int[nums.length + 1];
// calculate the cumulative sum of nums
for (int i = 1; i < preSum.length; i++) {
preSum[i] = preSum[i - 1] + nums[i - 1];
}
}
// query the cumulative sum of the closed interval [left, right]
public int sumRange(int left, int right) {
return preSum[right + 1] - preSum[left];
}
}class NumArray {
private:
// prefix sum array
vector<int> preSum;
public:
// input an array, construct the prefix sum
NumArray(vector<int>& nums) {
// preSum[0] = 0, for easy calculation of cumulative sum
preSum.resize(nums.size() + 1);
// calculate the cumulative sum of nums
for (int i = 1; i < preSum.size(); i++) {
preSum[i] = preSum[i - 1] + nums[i - 1];
}
}
// query the cumulative sum of the closed interval [left, right]
int sumRange(int left, int right) {
return preSum[right + 1] - preSum[left];
}
};class NumArray:
def __init__(self, nums: List[int]):
# prefix sum array
self.preSum = [0] * (len(nums) + 1)
# calculate the cumulative sum of nums
for i in range(1, len(self.preSum)):
self.preSum[i] = self.preSum[i - 1] + nums[i - 1]
def sumRange(self, left: int, right: int) -> int:
# query the cumulative sum of the closed interval [left, right]
return self.preSum[right + 1] - self.preSum[left]type NumArray struct {
// prefix sum array
preSum []int
}
// input an array, construct the prefix sum
func Constructor(nums []int) NumArray {
// preSum[0] = 0, for ease of calculating cumulative sum
n := len(nums)
preSum := make([]int, n + 1)
// calculate the cumulative sum of nums
for i := 1; i < n + 1; i++ {
preSum[i] = preSum[i - 1] + nums[i - 1]
}
return NumArray{preSum: preSum}
}
// query the cumulative sum of the closed interval [left, right]
func (this *NumArray) SumRange(left int, right int) int {
return this.preSum[right + 1] - this.preSum[left]
}var NumArray = function(nums) {
// prefix sum array
this.preSum = new Array(nums.length + 1);
// calculate the cumulative sum of nums
this.preSum[0] = 0;
for (let i = 1; i < this.preSum.length; i++) {
this.preSum[i] = this.preSum[i - 1] + nums[i - 1];
}
};
// query the cumulative sum of the closed interval [left, right]
NumArray.prototype.sumRange = function(left, right) {
return this.preSum[right + 1] - this.preSum[left];
};The core idea is to create a new array preSum, where preSum[i] records the cumulative sum of nums[0..i-1]. For example, see the image: 10 = 3 + 5 + 2.
Looking at this preSum array, if I want to find the sum of all elements in the index range [1, 4], I can get it by preSum[5] - preSum[1].
This way, the sumRange function only needs to perform one subtraction, avoiding the need for a for loop each time. The worst-case time complexity is constant O(1).
This technique is also widely used in real life. For instance, suppose you have several classmates, and each has a final exam score (out of 100). You need to create an API that, given any score range, returns how many students have scores within that range.
You can first use counting sort to calculate how many students have each specific score, and then use the prefix sum technique to implement the API for score range queries:
// stores all the students' scores
int[] scores;
// the full mark of the test paper is 100 points
int[] count = new int[100 + 1];
// record how many students have each score
for (int score : scores)
count[score]++;
// construct the prefix sum
for (int i = 1; i < count.length; i++)
count[i] = count[i] + count[i-1];
// use the prefix sum array 'count' to query score segmentsNext, let's see how the prefix sum concept can be applied in a two-dimensional array.
Prefix Sum in a 2D Matrix
This is LeetCode problem #304, "Range Sum Query 2D - Immutable". It's similar to the previous problem, where you were asked to calculate the sum of elements in a subarray. In this problem, you need to calculate the sum of elements in a submatrix within a 2D matrix:
304. Range Sum Query 2D - Immutable | 力扣 | LeetCode |
Given a 2D matrix matrix, handle multiple queries of the following type:
- Calculate the sum of the elements of
matrixinside the rectangle defined by its upper left corner(row1, col1)and lower right corner(row2, col2).
Implement the NumMatrix class:
NumMatrix(int[][] matrix)Initializes the object with the integer matrixmatrix.int sumRegion(int row1, int col1, int row2, int col2)Returns the sum of the elements ofmatrixinside the rectangle defined by its upper left corner(row1, col1)and lower right corner(row2, col2).
You must design an algorithm where sumRegion works on O(1) time complexity.
Example 1:
Input ["NumMatrix", "sumRegion", "sumRegion", "sumRegion"] [[[[3, 0, 1, 4, 2], [5, 6, 3, 2, 1], [1, 2, 0, 1, 5], [4, 1, 0, 1, 7], [1, 0, 3, 0, 5]]], [2, 1, 4, 3], [1, 1, 2, 2], [1, 2, 2, 4]] Output [null, 8, 11, 12] Explanation NumMatrix numMatrix = new NumMatrix([[3, 0, 1, 4, 2], [5, 6, 3, 2, 1], [1, 2, 0, 1, 5], [4, 1, 0, 1, 7], [1, 0, 3, 0, 5]]); numMatrix.sumRegion(2, 1, 4, 3); // return 8 (i.e sum of the red rectangle) numMatrix.sumRegion(1, 1, 2, 2); // return 11 (i.e sum of the green rectangle) numMatrix.sumRegion(1, 2, 2, 4); // return 12 (i.e sum of the blue rectangle)
Constraints:
m == matrix.lengthn == matrix[i].length1 <= m, n <= 200-104 <= matrix[i][j] <= 1040 <= row1 <= row2 < m0 <= col1 <= col2 < n- At most
104calls will be made tosumRegion.
Of course, you could use a nested for loop to traverse the matrix, but this would make the sumRegion function less efficient, and your algorithm would be less sophisticated.
Note that the sum of elements in any submatrix can be transformed into the sum of elements in several larger matrices surrounding it:
These four larger matrices share a common feature: their top-left corner is always the origin (0, 0).
A better approach for this problem is very similar to the prefix sum in a one-dimensional array. We can maintain a 2D preSum array that records the sum of elements in matrices with the origin as the top-left corner. This allows us to calculate the sum of any submatrix using just a few addition and subtraction operations:
class NumMatrix {
// define: preSum[i][j] records the sum of elements in the submatrix [0, 0, i-1, j-1] of matrix
private int[][] preSum;
public NumMatrix(int[][] matrix) {
int m = matrix.length, n = matrix[0].length;
if (m == 0 || n == 0) return;
// construct the prefix sum matrix
preSum = new int[m + 1][n + 1];
for (int i = 1; i <= m; i++) {
for (int j = 1; j <= n; j++) {
// calculate the sum of elements for each matrix [0, 0, i, j]
preSum[i][j] = preSum[i-1][j] + preSum[i][j-1] + matrix[i - 1][j - 1] - preSum[i-1][j-1];
}
}
}
// calculate the sum of elements in the submatrix [x1, y1, x2, y2]
public int sumRegion(int x1, int y1, int x2, int y2) {
// the sum of the target matrix is obtained by operating the four adjacent matrices
return preSum[x2+1][y2+1] - preSum[x1][y2+1] - preSum[x2+1][y1] + preSum[x1][y1];
}
}class NumMatrix {
private:
// define: preSum[i][j] records the sum of elements in the submatrix [0, 0, i-1, j-1] of matrix
vector<vector<int>> preSum;
public:
NumMatrix(vector<vector<int>>& matrix) {
int m = matrix.size(), n = matrix[0].size();
if (m == 0 || n == 0) return;
// construct the prefix sum matrix
preSum = vector<vector<int>>(m + 1, vector<int>(n + 1));
for (int i = 1; i <= m; i++) {
for (int j = 1; j <= n; j++) {
// calculate the sum of elements for each matrix [0, 0, i, j]
preSum[i][j] = preSum[i-1][j] + preSum[i][j-1] + matrix[i - 1][j - 1] - preSum[i-1][j-1];
}
}
}
// calculate the sum of elements of submatrix [x1, y1, x2, y2]
int sumRegion(int x1, int y1, int x2, int y2) {
// the sum of the target matrix is obtained by operating four adjacent matrices
return preSum[x2+1][y2+1] - preSum[x1][y2+1] - preSum[x2+1][y1] + preSum[x1][y1];
}
};class NumMatrix:
# Define: preSum[i][j] records the sum of elements in the submatrix [0, 0, i-1, j-1] of matrix
def __init__(self, matrix: List[List[int]]):
m, n = len(matrix), len(matrix[0])
if m == 0 or n == 0: return
# Construct the prefix sum matrix
self.preSum = [[0 for _ in range(n+1)] for _ in range(m+1)]
for i in range(1, m+1):
for j in range(1, n+1):
# Calculate the sum of elements for each matrix [0, 0, i, j]
self.preSum[i][j] = self.preSum[i-1][j] + self.preSum[i][j-1] + matrix[i-1][j-1] - self.preSum[i-1][j-1]
# Calculate the sum of elements in submatrix [x1, y1, x2, y2]
def sumRegion(self, x1: int, y1: int, x2: int, y2: int) -> int:
# The sum of the target matrix is obtained by operating on four adjacent matrices
return self.preSum[x2+1][y2+1] - self.preSum[x1][y2+1] - self.preSum[x2+1][y1] + self.preSum[x1][y1]type NumMatrix struct {
// define: preSum[i][j] records the sum of elements in the submatrix [0, 0, i-1, j-1] of matrix
preSum [][]int
}
func Constructor(matrix [][]int) NumMatrix {
m, n := len(matrix), len(matrix[0])
if m == 0 || n == 0 {
return NumMatrix{}
}
// construct the prefix sum matrix
preSum := make([][]int, m+1)
for i := 0; i <= m; i++ {
preSum[i] = make([]int, n+1)
}
for i := 1; i <= m; i++ {
for j := 1; j <= n; j++ {
// calculate the sum of elements for each matrix [0, 0, i, j]
preSum[i][j] = preSum[i-1][j] + preSum[i][j-1] + matrix[i-1][j-1] - preSum[i-1][j-1]
}
}
return NumMatrix{preSum: preSum}
}
// calculate the sum of elements in submatrix [x1, y1, x2, y2]
func (this *NumMatrix) SumRegion(x1 int, y1 int, x2 int, y2 int) int {
// the sum of the target matrix is obtained by operating on four adjacent matrices
return this.preSum[x2+1][y2+1] - this.preSum[x1][y2+1] - this.preSum[x2+1][y1] + this.preSum[x1][y1]
}var NumMatrix = function(matrix) {
var m = matrix.length, n = matrix[0].length;
if (m === 0 || n === 0) return;
// construct the prefix sum matrix
this.preSum = new Array(m + 1).fill(null).map(() => new Array(n + 1).fill(0));
for (var i = 1; i <= m; i++) {
for (var j = 1; j <= n; j++) {
// calculate the sum of elements for each matrix [0, 0, i, j]
this.preSum[i][j] = this.preSum[i-1][j] + this.preSum[i][j-1] + matrix[i - 1][j - 1] - this.preSum[i-1][j-1];
}
}
};
// calculate the sum of elements in submatrix [x1, y1, x2, y2]
NumMatrix.prototype.sumRegion = function(x1, y1, x2, y2) {
// the sum of the target matrix is obtained by operating on four adjacent matrices
return this.preSum[x2+1][y2+1] - this.preSum[x1][y2+1] - this.preSum[x2+1][y1] + this.preSum[x1][y1];
};This way, the time complexity of the sumRegion function is also optimized to O(1) using the prefix sum technique, which is a typical "space-for-time" approach.
The prefix sum technique is covered here. It's one of those skills that seems easy once you get it, but can be tricky if you don't. In practice, it's important to develop flexibility in your thinking to quickly recognize when a problem can be solved using prefix sums.
Besides the basic usage illustrated in this article, prefix sum arrays are often combined with other data structures or algorithmic techniques. I will provide examples and explanations in Prefix Sum Technique Practice Problems.
引用本文的题目
安装 我的 Chrome 刷题插件 点开下列题目可直接查看解题思路：
