Prefix Sum Array Technique
Note
Now all the plugins has supported English. I'm still improving the website...
This article will resolve
LeetCode | Difficulty |
---|---|
303. Range Sum Query - Immutable | 🟢 |
304. Range Sum Query 2D - Immutable | 🟠 |
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
nums
between indicesleft
andright
inclusive 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 ofnums
between indicesleft
andright
inclusive (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] <= 105
0 <= left <= right < nums.length
- At most
104
calls 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 res
type 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 method can achieve the desired effect, but it is very inefficient because the sumRange
method is frequently called, and its worst-case time complexity is , where N
represents the length of the nums
array.
The optimal solution to this problem is to use the prefix sum technique, which reduces the time complexity of the sumRange
function to . Simply put, do not use a for loop within sumRange
. How do we achieve this?
Let's look directly 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]
, as shown in the diagram: 10 = 3 + 5 + 2.
Looking at this preSum
array, if I want to find the sum of all elements within the index range [1, 4]
, I can derive it with preSum[5] - preSum[1]
.
In this way, the sumRange
function only needs to perform a single subtraction operation, avoiding the need for a for loop each time, with the worst-case time complexity being constant .
This technique is also widely used in everyday life. For example, if there are several students in your class, each with a final exam score (out of 100), you need to implement an API that, given any score range, returns how many students scored within that range.
You can first use counting sort to calculate how many students achieved each specific score, and then use the prefix sum technique to implement the score range query API:
// 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 segments
Next, let's look at how the prefix sum concept is applied in a two-dimensional array.
Prefix Sum in a 2D Matrix
This is LeetCode problem 304 "Range Sum Query 2D - Immutable". It is similar to the previous problem, where you needed to calculate the sum of elements in a subarray. In this problem, you are asked to calculate the sum of elements in a submatrix of 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
matrix
inside 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 ofmatrix
inside 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.length
n == matrix[i].length
1 <= m, n <= 200
-104 <= matrix[i][j] <= 104
0 <= row1 <= row2 < m
0 <= col1 <= col2 < n
- At most
104
calls will be made tosumRegion
.
Of course, you could use a nested for loop to traverse the matrix, but that would result in a higher time complexity for the sumRegion
function, reducing the efficiency of your algorithm.
Notice that the sum of the elements in any submatrix can be converted into operations involving the sums of a few larger matrices:
These four larger matrices share a common characteristic: their top-left corner is the origin (0, 0)
.
A better approach to this problem, similar to the prefix sum in a one-dimensional array, is to maintain a two-dimensional preSum
array. This array records the sum of elements in matrices with the origin as the top-left vertex, allowing you to calculate the sum of any submatrix using a few addition and subtraction operations:
class NumMatrix {
// definition: 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:
// definition: 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:
# definition: 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 {
// definition: 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.
Related Problems
You can install my Chrome extension then open the link.