Prefix Sum

Prefix Sum

·

3 min read

Let's see the leetcode Find Pivot Index problem.

The problem is to find the leftmost pivot index of an array. The pivot index is the index where the sum of all the numbers strictly to the left of the index is equal to the sum of all the numbers strictly to the index's right.

For example, take array [1,7,3,6,5,6] as an example, the pivot index is 3, because, 1+7+3 === 5+6. Note that the pivot itself is not included.

This is a simple problem. All we need to do is iterate the array and get the sum of left and right, and then compare them.

function pivotIndex(nums: number[]): number {
    for (let i = 0; i < nums.length; i++) {
        let left = 0;
        for (let j = 0; j < i; j++) {
            left += nums[j];
        }

        let right = 0;
        for (let k = nums.length - 1; k > i; k--) {
            right += nums[k];
        }

        if (left === right) return i;
    }
    return -1;
};

As you can see, we need to do another 2 for loop inside the outer for loop. So time complexity is O(n^2).

Can we do better?

This is the time to know the comcept of prefix sum. Prefix sum is kind of cumulative sum of all the element seen while iterating an array.

Take [1,7,3,6,5,6] as an example, the prefix sum of it can be expressed as below.

[1,7,3,6,5,6]

prefix sum =>

[
    1,
    1+7,
    1+7+3,
    1+7+3+6,
    1+7+3+6+5,
    1+7+3+6+5+6,
]

As you can see, this process is very similiar to the process to get the pivot's left sum and right sum.

Take [1,7,3,6,5,6] as an example, we can get the prefix sum and suffix sum first, and then use it to get the pivot index.

[1,7,3,6,5,6]

prefix sum =>

[
    0,
    1,
    1+7,
    1+7+3,
    1+7+3+6,
    1+7+3+6+5,
]

suffix sum => 

[
    7+3+6+5+6,
    3+6+5+6,
    6+5+6,
    5+6,
    6,
    0,
]

Note that because the pivot's value is not included. so the above prefix sum is adjusted accordingly.

With these two array, we know the left sum and right sum for every index. All we need to do is to compare them.

function pivotIndex(nums: number[]): number {

    const left: number[] = [];
    const right: number[] = [];

    // one loop to create prefix sum array and suffix sum array
    for (let i = 0; i < nums.length; i++) {
        const li = i;
        const ri = nums.length - i - 1;

        if (i === 0) {
            left[li] = 0;
            right[ri] = 0;
        } else {
            left[li] = left[li - 1] + nums[li - 1];
            right[ri] = right[ri + 1] + nums[ri + 1];
        }
    }

    // another loop to compare
    for (let i = 0; i < nums.length; i++) {
        if (left[i] === right[i]) {
            return i;
        }
    }

    return -1;
}

So with prefix sum, we make the complexity become O(n).

But as you can see, we need to use two array to store the prefix and suffix sum resutls. So space complexity is also O(n).

Actually, in this problem, we can just use 2 varibable to store the current prefix sum and suffix sum. And updating and comparing these 2 variables while looping.

function pivotIndex(nums: number[]): number {
    // suffix sum start from sum of all elements
    let right = nums.reduce((pre, cur) => pre + cur);
    // prefix sum start from 0
    let left = 0;

    /**
     * While iterating,
     * suffix sum delete left element one at a time
     * then, this is the corrent comparing time
     * prefix sum add left element one at a time
     */
    for (let i = 0; i < nums.length; i++) {
        right -= nums[i];
        if (left === right) return i;
        left += nums[i];
    }

    return -1;
};

We this small trick, the space comlexity becomes O(1).