Find the pivot index in Python

In this Answer, we’ll cover three methods to find the pivot index in Python. We’ll begin with the brute-force implementation and move to more efficient approaches afterward.

In this problem, we have to find the index of an array in which the sum of elements to the left and right of the index are equal. If no such index exists, we’ll return -1. This problem is illustrated in the image below:

An example of a pivot index
An example of a pivot index

Note: When all the elements succeeding it add up to 0, index 0 is the pivot index. When all the elements preceding it add up to 0, the last index is the pivot index.

Here’s a summary of each implementation:

  • Brute-force: The simplest approach that uses nested loops to check every index in the list and checks for a pivot. It calculates the sum of elements to the left and right of each index and checks if they’re equal. If they are, the pivot index has been found.

  • Prefix sums: This approach optimizes the solution by precomputing the list sum, which means we only iterate through the list once. It will compute the sum to the right of the index by subtracting the left sum and the current element’s value from the total sum.

  • Two-pointer: This approach optimizes the solution using two pointers, one from the left and the other from the right. It computes the total sum of the list and then initializes the left pointer to 0 and the right pointer to the sum. It then iterates through the list once, updating the right pointer by subtracting the current element’s value and checking whether the pivot index is found.

Brute-force implementation

Let’s take a look at the brute-force implementation of this problem:

def bruteforce(data):
  for i in range(len(data)):
    left = sum(data[:i])
    right = sum(data[i+1:])
    if left == right:
        return i
  return -1
data = [1, 7, 3, 6, 5, 6]
pivot = bruteforce(data)
if pivot != -1:
  print(f"Pivot index found {pivot}")
else:
  print("There is no pivot index")

This code is explained below:

  • Line 1: We define the bruteforce function, which takes data (list of numbers) as a parameter.

  • Line 2: We iterate through the list.

  • Line 3: We calculate the sum of the elements to the left side of the current index i.

  • Line 4: We calculate the sum of the elements to the right side of the current index i.

  • Line 5: If this condition is met, it means that i is the pivot index.

  • Lines 9–14: The main function to call our function bruteforce().

Optimized prefix sums implementation

Let’s take a look at the optimized prefix sums implementation:

def prefix(data):
  total = sum(data)
  left = 0
  for i in range(len(data)):
    right = total - left - data[i]
    if left == right:
      return i
    left += data[i]
  return -1
data = [1, 7, 3, 6, 5, 6]
pivot = prefix(data)
if pivot != -1:
  print(f"Pivot index found {pivot}")
else:
  print("There is no pivot index")

The code can be explained as follows:

  • Line 1: We define the prefix function, which takes data (list of numbers) as a parameter.

  • Line 2: We take the sum of the list and store it in total.

  • Line 3: We initialize the variable left to 0. This keeps track of the sum of elements to the left of the index while we iterate through the list.

  • Line 4: We iterate through the list. By using this for loop with range(len(data)), the code processes each element in the data list, allowing it to calculate the left and right sums and determine if a pivot index exists.

  • Line 5: We calculate the sum values toward the right of the index.

  • Lines 6–7: We compare the values of left and right sum. If this condition is met, it means that i is the pivot index.

  • Line 9: If no pivot index is found, we return -1.

  • Lines 12–17: The main function to call our function prefix().

Two-pointer implementation

Let’s take a look at the two-pointer implementation:

def pointer(data):
    left = 0
    right = sum(data)
    
    for i in range(len(data)):
        right -= data[i]
        if left == right:
            return i
        left += data[i]
    return -1
data = [1, 7, 3, 6, 5, 6]
pivot = pointer(data)
if pivot != -1:
    print(f"Pivot index found {pivot}")
else:
    print("There is no pivot index")

The code can be explained as follows:

  • Line 1: We define the pointer function, which takes data (list of numbers) as a parameter.

  • Lines 2–3: We initialize left to 0 and right to the sum of the elements in data.

  • Line 5: We iterate through the list using the same for loop in the previous implementation.

  • Lines 6–8: We subtract the value in data[i] from the right variable. Afterward, we’ll check if left is equal to right and if they are, then i is the pivot index.

  • Line 9: If they aren’t equal, we update the left value by adding the next element to the list.

Quiz

1

What is the purpose of finding a pivot index in an array?

A)

To sort the array efficiently.

B)

To find an index where elements to the left and right have equal sums.

C)

To find the maximum value in the array.

Question 1 of 30 attempted

Free Resources

Copyright ©2025 Educative, Inc. All rights reserved