The Prefix Sum method in Arrays
I was trying to solve the following problem a while back. It goes like this. Given an array and a set of queries, return the maximum sum in the final array after performing those queries. Let us understand with an example-
Given an array of n elements — All prefilled with 0. And another 2D array with a set of queries of the form —
a b k
1 5 3
4 8 7
6 9 1
The queries have to be performed on the array. So say 1 5 3
— this means you have to add 3 from index 1 to 5. Likewise, you have to do for each row of operations. Finally, you have to get the biggest number in the array.
Now, this is a fairly simple question. All you have to do is loop from indexes a
to b
and keep adding k
. Agreed, and that is what blunder I also made. Turns out that solution starts giving a timeout when the constraints are in the range 2*10^5
function arrayManipulation(n, queries) {
const arr = Array(n).fill(0);
let maximumSum = 0;
for (const query of queries) {
const a = query[0] - 1;
const b = query[1] - 1;
const k = query[2];
for (let i = a; i <= b; i++) {
arr[i] += k;
}
}
for (const element of arr) {
if (element > maximumSum) {
maximumSum = element;
}
}
return maximumSum;
}
This is the naive version. It worked fine till the constraints kicked in. That is when I came across something called the prefix-sum method. It’s tricky but really worth it. Here is the code —
function arrayManipulation(n, queries) {
const arr = Array(n).fill(0);
let max = 0;
let current = 0;
for (const query of queries) {
const [a, b, k] = query;
arr[a - 1] += k;
arr[b] -= k;
}
for (const element of arr) {
current += element;
if (current > max) {
max = current;
}
}
return max;
}
The idea is to add elements to the array in such a way that the cumulative sum is what is desired after each operation. Let us understand with an example.
arr = [0 0 0 0 0 0 0 0 0 0]
operation = 1 5 3
arr = [0 3 3 3 3 0 0 0 0 0] - // basic version
arr = [3 0 0 0 0 -3 0 0 0 0] - // prefix sum version
Now what just happened? The idea is to add k
to position a-1
and add -k
to position b
so that the cumulative sum becomes as desired. Now what is the cumulative sum? It is basically adding the previous element to the current element. Let us see what the cumulative sum becomes —
arr = [3 0 0 0 0 -3 0 0 0 0] - // prefix sum version
arr = [3 3 3 3 3 0 0 0 0 0] - // cumulative sum
arr = [3 3 3 3 3 0 0 0 0 0] - // basic version
You will see the cumulative and basic versions match. Let us try one more operation.
operation = 4 8 7
arr = [3 3 3 3 3 0 0 0 0 0] - // basic version before
arr = [3 3 3 10 10 7 7 7 0 0] - // basic version after
arr = [3 0 0 0 0 -3 0 0 0 0] - // prefix sum version before
arr = [3 0 0 7 0 -3 0 0 -7 0] - // prefix sum version after
arr = [3 3 3 10 10 7 7 7 0 0] - // cumulative sum version
The benefit is we were able to avoid the inner loop that runs from range a to b for each operation.
Just thought you would like to know there are crazy ways to do things.