LeetCode - Find Pivot Index

Author: Qi Wang

Edit This Page

Time Complexity: O(N)\mathcal O(N).

Solution 1

The pivot index must have the same sum on its left and its right. To calculate this, we can create a prefix sum of the nums\texttt{nums} array (in my solution it is cushioned with one 00 at index 00). Loop through every index in the psums\texttt{psums} array and check whether

psums[i]=?psums[size(nums)]−psums[i+1].\texttt{psums}[i] \stackrel{?}{=} \texttt{psums}[\text{size}(\texttt{nums})] - \texttt{psums}[i+1].

If this is true, we can return our current index because the problem wants the left-most index. If it finishes the for-loop without returning, it means there is no valid pivot index so return −1-1.

C++

Implementation

class Solution {
  public:
	int pivotIndex(vector<int> &nums) {
		vector<int> psums{0};
		// Constructing psums
		for (int num : nums) psums.push_back(psums.back() + num);


		for (int i = 0; i < psums.size() - 1; i++) {
			// Checking whether the left and the right are equal in area
			if (psums[i] == psums[nums.size()] - psums[i + 1]) { return i; }
		}
		return -1;
	}
};

Java

Implementation

class Solution {
	public int pivotIndex(int[] nums) {
		List<Integer> psums = new ArrayList<>();


		// Constructing psums
		psums.add(0);
		for (int num : nums) psums.add(psums.get(psums.size() - 1) + num);


		for (int i = 0; i < psums.length - 1; i++) {
			// Checking whether the left and the right are equal in area
			if (psums.get(i) == psums.get(nums.length) - psums.get(i + 1)) { return i; }
		}
		return -1;
	}
}

Python

Implementation

from typing import List




class Solution:
	def pivotIndex(self, nums: list[int]) -> int:
		psums = [0]
		for i in range(len(nums)):
			# Notice i in psums is always one less than the current index
			psums.append(psums[i] + nums[i])

Solution 2

The alternate solution does not require an array to store the sums. Instead, we only need the total sum of the array and the current sum at each index (exclusive of index). With this we check whether

curSum=?total−curSum−nums[i].\texttt{curSum} \stackrel{?}{=} \texttt{total} - \texttt{curSum} - \texttt{nums}[i].

Note: This is technically the same solution as solution 1, but it is another way to think about approaching this question.

C++

Implementation

class Solution {
  public:
	int pivotIndex(vector<int> &nums) {
		int total = 0;
		for (int x : nums) total += x;


		int curSum = 0;
		for (int i = 0; i < nums.size(); i++) {
			// curSum is up to i not inclusive of i.
			if (curSum == total - curSum - nums[i]) { return i; }

Java

Implementation

class Solution {
	public int pivotIndex(int[] nums) {
		int total = 0;
		for (int num : nums) total += num;


		int curSum = 0;
		for (int i = 0; i < nums.length; i++) {
			// curSum is up to i not inclusive of i.
			if (curSum == total - curSum - nums[i]) { return i; }
			// updating curSum
			curSum += nums[i];
		}
		return -1;
	}
}

Python

Implementation

from typing import List




class Solution:
	def pivotIndex(self, nums: List[int]) -> int:
		total = sum(nums)


		curSum = 0
		for i in range(len(nums)):
			# curSum is up to i not inclusive of i.
			if curSum == total - curSum - nums[i]:
				return i
			# updating curSum
			curSum += nums[i]
		return -1

Join the USACO Forum!

Stuck on a problem, or don't understand a module? Join the USACO Forum and get help from other competitive programmers!

Join Forum