Array: Prefix sum

Solving array problems can sometimes be very tedious and TLE for even simple problems is annoying. You might have faced this as well, right? I spent quite some time figuring out the optimal solution for the problem statements, then I came across some of the techniques which are intuitive and will save you a lot of time. After learning these techniques your perspective will change as to how to tackle a problem statement. I will be making a series of techniques in Array starting with Prefix sum so let's get started.

Problem Statement

Given an array A of size N and Q queries, where each query will have s and e.
You need to print the sum of all the elements in a given range (s, e).

Ex: 
      A = [3, 6, -1, 9, 2, -5, 8, 7]
      Q = [[1,4],[0, 7], [3, 3]]

      Q1: s = 1  e = 4
      Ans = 16
      Q2: s = 0  e = 7
      Ans = 29
      Q3: s = 3  e = 3
      Ans = 9

Brute force approach

Looking at the problem statement the simplest way to find the sum of all the elements in the given range would be to iterate from s to e and add up all the elements in the range.

function sum(A, Q) {
    for(let i = 0; i<Q.length; i++) {
        let sum = 0
        let s = Q[i][0]
        let e = Q[i][1]
        for(let j = s; j<=e; j++) {
            sum += A[j]
        }
        console.log(sum)
    }
}

The above solution looks fine but let's see if it can handle large input. In the worst case, the inner loop has to iterate N times if the range is from [0, N-1] and we can have a large number of Q queries, so the time complexity for the problem would be O(Q x N).

For an input of array A with 10⁵ elements and queries Q of length 10⁵, this loop will take 10¹⁰ iterations which might not be viable for the system, so let's optimize it!

Optimization

Before we get into the solution let us see an example to understand the approach better.

Let's say you have scores for a cricket match and you need to find the total score in a given range of overs:

Overs   1        2         3         4         5          6        7
Runs    3        15        33        40        45        57       74

You need to find the scores for the following overs:

Overs = 2 - 4
Run = 40 - 3 
    = 37

Overs = 3 - 7
Run = 74 - 15
    = 59

Overs = 6 - 6
Run = 57 - 45
    = 12

The above example is nothing but the summation of all the runs in a given range which is very much similar to our problem statement. This approach seems much simpler than iterating over all the elements in the range so let's see how we can use this method.

To get the prefix sum all we need to do is sum all the elements from start to the given index, let's get the prefix sum for the following array:

A = [3, 6, -1, 9, 2, -5, 8, 7]
PrefixSum = [3, 6, 8, 17, 19, 14, 22, 29]

Let's try to compute the queries to see if we get the correct result

Q1: s = 1   e = 4    
Ans = 19 - 3 = 13

Q2: s = 0       e = 7
Ans = 29

Q3: s = 3        e = 3
Ans = 17 - 8 = 9

Everything looks fine so let's jump to optimize the solution. First, let's get the prefix sum:

let ps = [] // prefix sum array
for(let i = 0; i<A.length; i++) {
    let sum = 0
    for(let j = 0; j<=i; j++) {
        sum = sum + A[j]
    }
    ps.push(sum)
}

We got prefix sum but there is still scope for improvement so let's see what we can improve, if you observe carefully we can use the sum computed for i - 1 element to compute i^th element. Let's update the code.

let ps = [] // prefix sum array
let sum = 0
for(let i = 0; i<A.length; i++) {
    sum = sum + A[i]
    ps.push(sum)
}

This is much better, now let's use the prefix sum to solve the problem statement.

function sum(A, Q) {
    let ps = [] // prefix sum array
    let sum = 0
    for(let i = 0; i<A.length; i++) {
        sum = sum + A[i]
        ps.push(sum)
    }

    for(let i = 0; i<Q.length; i++) {
        let s = Q[i][0]
        let e = Q[i][1]
        let sum = ps[e] - ps[s - 1]
        console.log(sum)
    }
}

This is much better, we got rid of the nested loop, but there are some edge cases where this might fail so let's look into it. When the value of s = 0 we will get out of bound error since we will be trying to access the index -1 so let's add a condition for that.

function sum(A, Q) {
    let ps = [] // prefix sum array
    let sum = 0
    for(let i = 0; i<A.length; i++) {
        sum = sum + A[i]
        ps.push(sum)
    }

    for(let i = 0; i<Q.length; i++) {
        let s = Q[i][0]
        let e = Q[i][1]
        if(s === 0) {
            let sum = ps[e]
            console.log(sum)
        } else {
            let sum = ps[e] - ps[s - 1]
            console.log(sum)
        }
    }
}

This is much better than we had before. Let's see the time complexity for this solution.

To compute prefix sum we will iterate N times which is O(N) and for resolving the query we will iterate Q times which is O(Q), now computing sum is a constant operation so it will be O(1). So the overall time complexity for this solution will be O(N + Q) which is much better than O(N x Q).

For an input of array A with 10⁵ elements and queries Q of length 10⁵, this solution will take 2 x 10⁵ iterations which is much fast.

Hope this was helpful, feel free to ask any doubt in the comment section and stay tuned for more array techniques.