LeetCode - Find Pivot Index

Author: Qi Wang

Appears In

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

//Created by Qi Wang
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++){

Java

Implementation

//Created by Qi Wang
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);

Python

Implementation

# Created by Qi Wang
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]);
		
		for i in range(len(nums)):
			# Checking whether the left and the right are equal in area
			if psums[i] == psums[len(nums)] - psums[i+1]:
				return i;
			
		return -1;

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

//Created by Qi Wang
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++){

Java

Implementation

//Created by Qi Wang
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.

Python

Implementation

# Created by Qi Wang
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