[Problem Solving] Find Pivot Index

Solution #1: To get started with a brute force approach, let’s consider each number as a pivot and calculate the left and right side sums until they equal each other.

Pivot = 1, pivotIndex = 0, leftSum = 0, rightSum = 27

Pivot = 7, pivotIndex = 1, leftSum = 1, rightSum = 20

Pivot = 3, pivotIndex = 2, leftSum = 8, rightSum = 17

Pivot = 6, pivotIndex = 3, leftSum = 11, rightSum = 11. Voila!! Here leftSum = rightSum, so return 3 as the answer.

Time Complexity: A worst case scenario is when we have to traverse the entire array and the pivot index is the last index (n-1) of the array which will be O(N). Calculating left and right side sums adds another O(N) and so it will be an O(N²) algorithm. To improve our algorithm efficiency, we have to try to do this in linear time. In the previous approach, calculation of leftSum added to increased complexity, how can we optimize this part? What if we had access to the leftSum and rightSum values in advance?

Solution #2: Let’s pre-calculate the leftSum and rightSum values for each pivot and store these values in two arrays.

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

leftSums = [0, 1, 8, 11, 17, 22, 28]

rightSums = [27, 20, 17, 11, 6, 0]

In leftSums and rightSums, we need to find the index whose values are equal, in this case it’s 11 ie., leftSums[3] = rightSums[3].

Time Complexity: In Solution #2, our time complexity will go down from O(N²) to O(N) since we use pre-calculated sum values to find the pivot index ie., O(N) for pre-calculation plus another O(N) while traversing to find pivot index .

Space Complexity: Using two extra arrays to pre-compute values adds additional 2N space. So the overall space complexity will be O(N).

Code for Solution #2:

def pivotIndex(self, nums: List[int]) -> int:
        
        if not nums:
            return -1
        
        n = len(nums)
        leftSums, rightSums = [0] * n, [0] * n
       
        for i in range(1, n):
            leftSums[i] = leftSums[i-1] + nums[i-1]
            
        for i in range(n-2, -1, -1):
            rightSums[i] = rightSums[i+1] + nums[i+1]
            
        for i in range(n):
            if leftSums[i] == rightSums[i]:
                return i
            
        return -1

Slight improvement on space complexity:

Anytime we are traversing an array both ways and then iterating it all over again, it’s possible to eliminate one of the traversals and calculate values on the fly. In the above code snippet, the following chunks of code are kind of doing the same thing, let’s try to avoid code duplication.

for i in range(1, n):
    leftSums[i] = leftSums[i-1] + nums[i-1]

for i in range(n):
    if leftSums[i] == rightSums[i]:
        return i

We can replace the above two blocks by calculating leftSum on the fly as shown below.

for i in range(0, n):

    if i > 0:
        leftSum += nums[i-1]
    if leftSum == rightSums[i]:
        return i

Time Complexity: O(N)

Space Complexity: O(N)

Testing: For array problems like these, it’s important to test the following cases.

• Empty array: [] => returns -1
• Single value array: [1] => returns 0
• Double value array: [1, 2] => returns -1
• A random array: [2, 3, 4, 5] => returns 2
• All given test cases in the problem statement
• Arrays that have one pivot index: [1, 7, 3, 6, 5, 6] => returns 3
• Arrays that have more than one pivot index: [1, 0, 1, 0] => returns 0
• Arrays with no pivot index: [1, 2, 3, 4] => returns -1

Can we do better? Yes!

Solution #3: Introducing prefix sum

If we are given the leftSum of a pivot, what else do I need to get the rightSum? The total sum of the array.

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

totalSum = 1 + 7 + 3 + 6 + 5 + 6 = 28

pivot = 7, pivotIndex = 1, leftSum = 1, rightSum = 28 – 1 – 7 = 20

pivot = 6, pivotIndex = 3, leftSum = 11, rightSum = 28 – 11 – 6 = 11

To generalize, rightSum(i) = totalSum – leftSum – nums[i]. Now let’s alter the code accordingly.

def pivotIndex(self, nums: List[int]) -> int:
        
        if not nums:
            return -1
        
        n = len(nums)
        leftSum = 0
        totalSum = sum(nums)
        
        for i in range(n):
            if leftSum == totalSum - leftSum - nums[i]:
                return i
            
            leftSum += nums[i]
            
        return -1

Time Complexity: O(N) for calculation sum of nums, O(N) for finding pivot index, so the overall time complexity will be O(N).

Space Complexity: O(1) since we use constant space to store totalSum and leftSum values.

Want to solve this problem? https://leetcode.com/problems/find-pivot-index/

Hope you enjoyed this step by step approach to iteratively improve upon a solution for a problem. Have a great day!