Posted in Arrays and Strings, coding interview preparation, Data Structures/ Leetcode, Leetcode Top 150, Two Pointer Technique

[Problem Solving] Remove duplicates in a sorted array in-place

Problem Statement

Given an array of integers sorted in increasing order, remove the duplicates in-place such that each unique element appears only once. The relative order of the elements should be kept the same. Return k after placing the final result in the first k slots of nums.

Examples

Input: nums = [1,1,2]
Output: 2, nums = [1,2,_]
Input: nums = [0,0,1,1,1,2,2,3,3,4]
Output: 5, nums = [0,1,2,3,4,_,_,_,_,_]

This is one of my favorite and easy array problems to start thinking in terms of two pointers to solve array problems. Why two pointers in this case?

We need a pointer that will read the elements of the array one by one. We can use another pointer to overwrite the elements in place. Also whenever we want to modify elements of an array in place – think about whether the two pointer technique can be applied!

The read pointer, let’s say r is pretty straightforward where we increment its value by one during each iteration. The only problem to solve here is to understand when to move the write w pointer.

In this problem, we are going to always compare two elements to each other to find out if it is a repeated value. This naturally leads to maintaining two variables prev and curr to keep track of the previous element to compare the current one with. Now for moving the write pointer,

  • Each time prev = curr, we move the read pointer r and not the write pointer w. This is so that we keep reading the elements until we find the element to overwrite with.
  • When prev != curr, we can now overwrite the array element, increment w and r as well as set prev to curr.

Let’s see how this works using the following diagram. Note that only r and w are variables that we use for the two pointer technique here. Prev and curr are variables to keep track of the previous and current elements.

Complexity

O(N) to traverse the array and O(1) space since we modify the array in-place.

Solution

def removeDuplicates(self, nums: List[int]) -> int:
    if not nums:
        return 0
    if len(nums) == 1:
        return 1
             
    r, w = 1, 1
    prev = nums[0]
    while (r < len(nums) and w < len(nums)):
        curr = nums[r]
        if prev != curr:
            nums[w] = curr
            w += 1
            prev = curr
        r += 1
    return w

Hope you enjoyed learning to solve this array problem with the two pointer technique.

Posted in Data Structures/ Leetcode

Monotonic Stack

Today let us learn about monotonic stacks and their usage.

A monotonic stack is one in which the elements of the stack are always increasing or decreasing. This ensures that the stack is always sorted in some order. Some interesting applications of monotonic stacks are to

  • Find the next smaller or greater element in an array
  • Find the previous smaller or greater element in an array
  • Find the minimum or maximum element in a sliding window etc.

Let’s take a look at an example problem below.

Find the next greater element in an array

The next greater of an element X in an array is the first greater element to the right side of X. For example:

[1, 2, 3, 4]

  • Next greater element (1) => 2
  • Next greater element (2) => 3
  • Next greater element (3) => 4

In this problem, we can use a monotonic stack to push elements into and keep popping the elements as long as the top most element has a value smaller than the current element whose next greater element we want to find. That’s a mouthful, so let’s look at an example!

We have the following array.

Input: [1, 3, 4, 2, 5] and here’s an array of the next first greater elements to the right of each element in the array.

Output: [3, 4, 5, 5, -1] (-1 when no next greater element exists)

Algorithm:

  • Iterate the given array from i = 0 to i = N – 1.
  • Push the element array[i] in the monotonic stack if either the stack is empty or the top of stack > array[i]. We are thereby building a monotonic stack decreasing from the bottom to top.
  • In this process, we have to pop the elements in stack if they are lesser than array[i] which is the first greater element found to the right of the popped one.

Look at how the stack and the mapping changes as we iterate the array in the following diagram.

Finally push 5.

Time Complexity: O(N) where N is the size of the array that is iterated. This is because we iterate the array only once and each element is pushed and popped exactly once.

Space Complexity: O(N) since the mapping and the stack will contain at most N elements.

Code:

In the following code nums1 contains the elements whose next greater element should be found in nums2.

class Solution:
    def nextGreaterElement(self, nums1: List[int], nums2: List[int]) -> List[int]:
        num_to_next_greater_mapping = collections.defaultdict(lambda: -1)
        mono_stack = []
        
        for index, val in enumerate(nums2):
            if mono_stack and val > mono_stack[-1]:
                while mono_stack and mono_stack[-1] < val:
                    num_to_next_greater_mapping[mono_stack[-1]] = val
                    mono_stack.pop(-1)
            mono_stack.append(val)
            
        results = []
        for n in nums1:
            results.append(num_to_next_greater_mapping[n])
            
        return results
            

Stay tuned for more of these problems in this space.

Posted in Arrays and Strings, Data Structures/ Leetcode, Recursion

Regular Expression Matching

Given an input string and a regex pattern, implement regular expression matching with support for ‘.’ and ‘*’.

  • ‘.’ matches any single character
  • ‘*’ matches zero or more of the preceding element (the char before *)

Examples

  1. s = ‘aa’, p = ‘ab’ Not a match
  2. s= ‘bbbbb’, p = ‘b*’ Match because ‘*’ matches zero or more of the preceding element b
  3. s= ‘fljflf’, p = ‘.*’ Match because we want to match zero or more of the preceding element which is any character in this case
  4. s = ‘oop’, p = ‘r*o*p’ Match because r is present 0 times, o is present twice and p is a direct match

Thought Process

Let’s approach this problem step by step.

Case 1: The simplest case would be when the given string and the pattern do not contain any ‘.’ or ‘*’ characters but only alphabets.

s = 'aa' and p = 'ab'

In this case we would check the string and the pattern from left to right to see if each character in the string matches the characters in the pattern.

a = a matched so far
b != a Match failed so return False

Case 2: When no ‘*’ characters are present but only ‘.’. In this case ‘.’ can match any single character. So we check for that condition and continue to check the remaining characters in the given string and the pattern using recursion.

def isRegexMatch(s, pattern):
    # If pattern is empty, as long as string is also empty => true
    # If pattern is empty but string is not empty => return false
    if not pattern:
        return not s

    # Check the first character match in s and pattern
    # String should not be empty
    # The first characters exactly match (Or)
    # pattern's first char contains "." to match any char in s
    is_first_char_match = len(s) > 0 and (s[0] == pattern[0] or pattern[0] == ".")
    
    # Continue checking the remaining characters using recursion
    return is_first_char_match and isRegexMatch(s[1:], pattern[1:])

Case 3: When the regex pattern contains both ‘.’ and ‘*’ characters, things get a little complicated.

Note that the pattern will never begin with ‘*’ because it denotes that a preceding character can be present 0 or more times. So it will ALWAYS be present only as the second character when that part of the pattern is being processed after a recursive call, For example: a*, .*

So we can have a condition that evaluates this logic as shown below.

if len(pattern) >= 2 and pattern[1] == "*"

If the above condition holds true, then we check regex match differently due to the presence of ‘*’. Let’s think about how we can do that. There are two valid scenarios that are allowed here.

Scenario 1: s = "aaaa", pattern = "c*a*"

In this case len(pattern) = 4 which is >= 2 and pattern[1] = ‘*‘. We see that c is not present in the string s but it can occur zero 0 or more times. In this case we can ignore checking the first two characters of the pattern and continue the recursion as shown below.

isRegexMatch(s, pattern[2:]) => isRegexMatch("aaaa", "a*")
Scenario 2: s = "aaaa", pattern = "a*"

In this case len(pattern) = 2 which is >= 2 and pattern[1] = ‘*‘. s contains ‘a’ which is a direct match with the first character in the pattern. So we check for first character match and continue recursion on the rest of the string.

is_first_char_match and isRegexMatch(s[1:], pattern) => a = a and isRegexMatch("aaa", "a*")

Note that we pass the full pattern string to the recursion call here because ‘*’ can match zero or more characters of ‘a’. So we want to check the string for that as well. The following recursive calls will hold true.

isRegexMatch("aaa", "a*") =>
is_first_char_match and isRegexMatch(s[1:], pattern) => a = a and isRegexMatch("aa", "a*")
isRegexMatch("aa", "a*") =>
is_first_char_match and isRegexMatch(s[1:], pattern) => a = a and isRegexMatch("a", "a*")
isRegexMatch("a", "a*") =>
is_first_char_match and isRegexMatch(s[1:], pattern) => a = a and isRegexMatch("", "a*")
isRegexMatch("", "a*") =>
isRegexMatch(s, pattern[2:]) => isRegexMatch("", "") so final result will be true.
if len(pattern) >= 2 and pattern[1] == "*":
    return isRegexMatch(s, pattern[2:]) or (is_first_char_match and isRegexMatch(s[1:], pattern))

Complexity

For a given string of length S and pattern P, we make recursive calls using indexes s[i:] and pattern[2j:] (When ‘*’ is present). The recursive tree can go to a depth of (S+P/2) which will have 2h leaf nodes where h = (S+P/2), so 2(S+P/2) leaf nodes. This means the recursion will branch into 2(S+P/2) sub-problems.

If we pass the entire string and pattern at the worst case, the complexity of solving a sub-problem will be O(S+P). So the time complexity will roughly be

O((S+P) * 2(S+P/2))

To read more about the complexity calculation for this problem, here’s a detailed article.

Putting it all together.

Solution

def isRegexMatch(s, pattern):
    # If pattern is empty, as long as string is also empty => true
    # If pattern is empty but string is not empty => return false
    if not pattern:
        return not s

    # Check the first character match in s and pattern
    # String should not be empty
    # The first characters exactly match (Or)
    # pattern's first char contains "." to match any char in s
    is_first_char_match = len(s) > 0 and (s[0] == pattern[0] or pattern[0] == ".")
    
    # When '*' is present in the second char of pattern
    if len(pattern) >= 2 and pattern[1] == "*":
        return isRegexMatch(s, pattern[2:]) or (is_first_char_match and isRegexMatch(s[1:], pattern))

    else:
        # Continue checking the remaining characters using recursion
        return is_first_char_match and isRegexMatch(s[1:], pattern[1:])

Hope you enjoyed solving the regular expression problem. In the next post, we will look at how to improve upon this solution using Dynamic Programming!

Posted in Arrays and Strings, Data Structures/ Leetcode, Two Pointer Technique

[Problem Solving] Three Sum – Using Two Sum

To understand the following problem, we will be using what we learnt in the two sum problem. If you haven’t read it yet, do make sure to understand it before proceeding. Here’s the link.

Problem Statement

We have an array of integers called nums. We need to return three numbers let’s call it nums[i] = x, nums[j] = y and nums[k] = z that add up to 0. The indices of the returned numbers should all be unique ie., i != j != k

Assumptions

  • Solution set should not contain duplicate triplets
  • Empty array results in an empty set
  • If there are less than three numbers in the input array, return an empty set

Examples

Input: nums = [-1,0,1,2,-1,-4]

Output: [[-1,-1,2],[-1,0,1]]

Thought process

The brute force way to do this would be to loop using three for loops to search for three numbers in the array that sum up to 0. But this would have an O(n3) time complexity. So we need to be slightly better than that. The next best in time is an O(n2) solution. Remember that we should not return duplicate indices.

Let’s say nums[i] + nums[j] + nums[k] = 0

Then nums[i] + nums[j] = – nums[k]

This now resembles a two sum problem where we need to look for nums[i] and nums[j] such that they add up to a target = – nums[k]

To condense the problem, what we can do is loop through the array from k = 0 to k = n-1 and for every nums[k], find two numbers nums[i] and nums[j] that add up to a target = – nums[k].

How can we avoid adding duplicate triplets to the results list? A naive solution is to add all the duplicates to the result list and in the end return the set(result) to remove duplicates. This way we still do unnecessary work on finding the same two numbers that add up to a target value -nums[k].

To avoid executing a two sum loop on duplicates, let’s reuse our solution from the two sum problem where the input array was sorted and contained duplicates. We can sort the nums list and skip searching for two sum values whenever the current nums[k] = nums[k-1]. Sorting makes it easier to club duplicate values together. For ex:

Input: [-1,0,1,2,-1,-4]

Sorted List: [-4, -1, -1, 0 ,1,2]

One of the solutions here that adds up to 0 will be [-1, 0,1]. Now if we did not skip for target = nums[2] = -1, we will again get a solution set of [-1, 0, 1] in the result.

Complexity

O(N2) time complexity and O(1) extra space

Solution

def three_sum(nums):
    """
    :type nums: List[int]
    :rtype: List[List[int]]
    """
    # Edge case when the array is empty
    if not nums:
       return []
    
    n = len(nums)
    # Another edge case when the input array does not have 3 elements
    if len(nums) < 3:
       return []
        
    results = []

    # Sorting helps to club duplicate values together
    nums.sort()
        
    for k in range(0, n):
        # If the current num is equal to the previous value, then skip execution because we will be getting a duplicate solution set
        if k > 0 and nums[k] == nums[k-1]:
            continue
                
            target = -1 * nums[k]
            start = k + 1
            end = n - 1
            
            two_sum(nums, target, results, start, end)
    return results

# Same as the two sum solution that we coded in the previous article
def two_sum(nums, target, results, s, e): 

    while s < e:
            
        if nums[s] + nums[e] == target:
            results.append([nums[s], nums[e], -target])
            s += 1 # Advance the starting pointer to continue to look for other solutions

            # Skip duplicate solutions in two sum
            while s < e and nums[s] == nums[s-1]:
                s += 1
            
        elif nums[s] + nums[e] < target:
           s += 1
        else:
           e -= 1

Posted in Arrays and Strings, Data Structures/ Leetcode, Two Pointer Technique

[Problem Solving] Two Sum – Unsorted and Sorted Array, With and Without Duplicates

Two Sum – Problem Statement

We have an array of integers called nums and a target value. We need to return the indices of two numbers, say x and y that add up to target.

Assumptions

  • Each given input will have exactly one solution
  • Same element cannot be used twice
  • Order of indices does not matter

Examples

Input: nums = [2, 7, 11, 15], target = 9

Output: [0, 1] where x = 2 and y = 7

Input: nums = [4, 4], target = 8

Output: [0, 1] where x = 4 and y = 4

Thought process

What’s the brute force way to check if two numbers add up to a target value in an array?

For each value nums[i] where i = 0 initially, we traverse the list from nums[i+1] and see which two indices add up to a target value. The complexity of this will be O(n2) since we loop through the remaining n-i elements for each i. So the first optimization here would be to try and see if we can reduce the traversal to just once.

For a value x, to check if there is another value y = target – x present in the array, we need a way to do a fast lookup without going through the whole array. A dictionary is the data structure that comes to mind in that case. How to design this data structure? What will our keys and values be?

To search for a value y = target – x, we have to store the values in nums as the keys in dictionary to search for. The value for each key can then be its index in the nums array so that we can directly return the indices as our result.

Complexity

O(N) time to traverse the array once, O(1) to look up values.

O(N) space for the look up dictionary

Solution

def two_sum(self, nums, target):
    """
    :type nums: List[int]
    :type target: int
    :rtype: List[int]
    """
        
    look_up = {}
        
    for i, x in enumerate(nums):
       y = target - x

       # Additional condition to make sure we don't return the same index twice
       if y in look_up and look_up[y] != i:
           return [i, look_up[y]]
           look_up[x] = i
            
     return []  

Two Sum with Sorted Array – Problem Statement

This is a slight variation of the problem above where the input array nums is sorted in ascending order and we need to find two numbers that add up to a target value.

Assumptions

  • Each input will have exactly one solution
  • Same element cannot be used twice to add up to a target

Examples

Input: nums[2, 7, 11, 15], target = 9

Output: [0, 1]

Thought Process

Again, the brute force way of solving this problem would be to traverse the array using two for loops to find two numbers that add up to a target. This way we are not using the added benefit that the array is sorted in ascending order.

This means we can apply the two pointer technique here. This works great in many scenarios especially when the input array is sorted. We can use two pointers starting at two ends of the array. If the numbers they are pointing to do not add up to our target value, there are two decisions that can be made from then on. If the sum is greater than target, this means we have to move the end pointer to the left, otherwise move the start pointer to the right. See below.

Two Pointer Technique

Complexity

O(N) time since we traverse the array at most once in case the two elements are at the center of the array.

O(1) space since no additional space is used except for a couple of variables.

Solution

def two_sum(nums, target):
        """
    :type nums: List[int]
    :type target: int
    :rtype: List[int]
    """
        
    s = 0
    e = len(nums) - 1
        
    while s < e:
            
        if nums[s] + nums[e] == target:
           return [s, e]
            
        elif nums[s] + nums[e] < target:
           s += 1
        else:
           e -= 1
                
     return []

Two Sum with Sorted Array and Duplicates

Let’s consider the exact same problem as above but now our result can have multiple solutions that add up to the same target value but it cannot contain duplicates.

Instead of returning the indices, let’s return the numbers to make it easier for understanding.

Assumptions

  • Each input will have multiple solutions
  • Same element cannot be used twice to add up to a target
  • The solution cannot contain duplicate sets

Examples

Input: nums = [2, 2, 4, 4, 5, 7, 7, 9, 11, 15], target = 9

Output: [[2, 7], [4, 5]]

Thought Process

In order to return multiple solutions, once we find nums[s] + nums[e] = target, we have to advance our start s pointer and continue to look for other solutions. We can collect the solutions inside a result list. If we don’t account for duplicate solutions, this is how our code will look. For the given input above, this code will return all solutions including the duplicates [[2, 7], [2, 7], [4, 5], [4, 5]].

def two_sum(nums, target):
    """
    :type nums: List[int]
    :type target: int
    :rtype: List[int]
    """
        
    s = 0
    e = len(nums) - 1
    result = []
    
    while s < e:
            
        if nums[s] + nums[e] == target:
            result.append([nums[s], nums[e]])
            s += 1 # Advance the starting pointer to continue to look for other solutions
            
        elif nums[s] + nums[e] < target:
           s += 1
        else:
           e -= 1
                
    return result

In order to avoid duplicate solutions since our array is already sorted, we can compare the current value with the previous value. If they are the same, we are about to build the exact same solution that was found in the previous run, so we can skip it. For example, in [2, 2, 4, 4, 5, 7, 7, 9, 11, 15], when i = 1, nums[i] = 2 = nums[i-1] = 2, which will again yield [2, 7] as a duplicate solution, so let’s skip it and advance the starting s pointer. Then our solution can be slightly modified as shown below.

Complexity

O(N) time

O(1) space since we don’t usually include result as additional space

Solution

def two_sum(nums, target):
    """
    :type nums: List[int]
    :type target: int
    :rtype: List[int]
    """
        
    s = 0
    e = len(nums) - 1
    result = []
    
    while s < e:
            
        if nums[s] + nums[e] == target:
            result.append([nums[s], nums[e]])
            s += 1 # Advance the starting pointer to continue to look for other solutions

            # Skip duplicate solutions
            while s < e and nums[s] == nums[s-1]:
                s += 1
            
        elif nums[s] + nums[e] < target:
           s += 1
        else:
           e -= 1
                
    return result

Posted in Data Structures/ Leetcode, June Leetcoding Challenge, Trees

[Problem Solving] Invert a Binary Tree

We are given a binary tree as shown below.

Invert a binary tree

Let’s observe the given tree and compare it with the output. What do we notice?

The left child of a node becomes the right child and vice versa. So we start at the root node and invert the left and right children.

     1
   /   \
  2     3

becomes

     1
   /   \
  3     2

We then invert the left and right sub-trees starting at each of those children. This leads to a naturally recursive approach.

     1
   /   \
  3     2
 / \   / \
6   7 4   5

becomes

     1
   /   \
  3     2
 / \   / \
7   6 5   4

An intuitive way to solve most tree problems is to use recursion. Now we have to decide a few things before coding.

Base case:

Here, the base case would be a NULL node. We simply return at that point since there are no nodes left to invert.

Recursive strategy:

Our recursive strategy involves the following steps.

  • Invert the left and right child of current node
  • Recursive call to invert the left sub-tree starting at the left child
  • Recursive call to invert the right sub-tree starting at the right child

Time Complexity: We will end up visiting each node in the tree, so the time complexity will be O(N) where N is the number of nodes in the tree.

Space Complexity: In recursion problems, an internal stack is used to temporarily hold elements, so the space complexity will be O(H) where H is the height of the tree. This is because as we recurse over left and right sub-trees, the function calls will go as deep as the tree height. In the worst case, H -> N, so the overall space complexity will be O(N) where N is the number of nodes in the tree.

Code for the above solution:

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

def invertTree(self, root: TreeNode) -> TreeNode:

    def invertTreeHelper(root):

        if not root:
            return

        root.left, root.right = root.right, root.left
        invertTreeHelper(root.left)
        invertTreeHelper(root.right)

    invertTreeHelper(root)
    return root

Testing:

These are some of the cases that are good to test for tree problems.

  • Given test case: [1,2,3,4,5,6,7]
  • Empty tree: []
  • No children: [1]
  • No right child: [1, 2]
  • Right skewed tree: [1,null,2,null,3,null,4,null,5,null,6,null,7]
  • Left skewed tree: [1,2,null,3,null,4,null,5,null,6,null,7]

Want to try to solve this problem? https://leetcode.com/problems/invert-binary-tree/

Hope you learnt how to invert a binary tree. Have a great day!

Posted in Arrays and Strings, Data Structures/ Leetcode

[Problem Solving] Find Pivot Index

Let’s say we are given an array of integers and we want to find a pivot index such that the sum of numbers to the left of this index equals the sum of numbers to the right.

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

1 + 7 + 3 = 5 + 6 = 11. In this example, the pivot value is 6, whose array index is 3. If no such index exists, return -1.

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(N2) 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(N2) 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!