Posted in Arrays and Strings, June Leetcoding Challenge, Python

5 Pythonic ways to reverse a string

Write a function that reverses a string. The input string is given as an array of characters char[].

Input: A = [“h”,”e”,”l”,”l”,”o”]
Output: [“o”,”l”,”l”,”e”,”h”]

Solution #1: Arrays ie., lists in Python are mutable data structures which means they can be modified in place without having to allocate extra space. For this problem, creating a new array using the input array is a trivial solution. So let’s directly dive into an O(1) memory solution.

In Python, we can refer to the elements in reverse order starting at the index -1.

A[-1] = “o”

A[-2] = “l”

A[-5] = “h”

This is quite helpful in manipulating array indexes. In order to reverse the elements, we don’t have to traverse the full list but swap the first half of the list with the second half starting from the end.

ie., swap A[i] with A[-(i+1)] for i = 0 to i = len(A) / 2

Time Complexity: We traverse only half the list while swapping elements with the second half of the list. So the time complexity will be O(N/2) which is O(N) overall.

Space Complexity: O(1)

Code for solution #1:

def reverseString(self, s: List[str]) -> None:
    """
    Do not return anything, modify s in-place instead.
    """
    for i in range(len(s) // 2):
        s[i], s[-(i+1)] = s[-(i+1)], s[i]

Solution #2: A Python in-built function that does exactly what we did above is reverse(). The list is modified in place by swapping the first half of the list with the second half the same way as shown above. Read more about reverse() here.

Code for solution #2:

def reverseString(self, s: List[str]) -> None:
    s.reverse()

Time Complexity: O(N)

Space Complexity: O(1)

Solution #3: If we are allowed to create and return a new list, then we can use Python reversed() function as shown below. reversed() returns an iterator in Python, so we have to wrap its result in a list(). Read more about it here.

Code for solution #3:

def reverseString(self, s: List[str]) -> List[str]:
    return list(reversed(s))

Time Complexity: O(N)

Space Complexity: O(N) since we are creating a new list.

What if we are given a string and not an array of characters?

An important distinction here is that strings are immutable, so we cannot leverage any in-place reversals here. Let’s look at some Pythonic ways of reversing a string.

Input: “hello”

Output: “olleh”

Solution #4: Using Python’s elegant index slicing.

Code for solution #4:

def reverseString(self, s: str) -> str:
    return s[::-1]

Okay, what was that? In Python there is a concept called slicing. The syntax for slicing is shown below.

A[start: stop: step]

  • start – The index to start slicing (default: 0)
  • stop – The index to stop slicing (default: len(s))
  • step – The increment between each index for slicing

Examples:

  • A[0:1] slices the first element
  • A[1:4] slices the elements from index 1 to 3

The -1 step in the code indicates that we will start at the end and stop at the start instead of the other way around, thereby reversing the string.

Read more about Python slicing here.

Time Complexity: O(N)

Space Complexity: O(N) since a new string is created.

Solution #5: Another way to improve the readability of your code instead of using the above syntax which a non-python developer may not immediately be able to comprehend is to use reversed() to reverse a string.

  • Use reversed() to create an iterator of the string characters in the reverse order.
  • Use join to merge all the characters returned by the iterator in the reverse order.
def reverseString(self, s: str) -> str:
    return "".join(reversed(s))

Time Complexity: O(N)

Space Complexity: O(N) since a new string is created.

Let’s summarize all the different ways.

InputMethodTime ComplexitySpace Complexity
list[str]Swap elements from first half with second half of the array.O(N)O(1)
list[str]Built-in list.reverse()O(N)O(1)
list[str]list(reversed(s))O(N)O(N)
strSlicing str[: : -1]O(N)O(N)
str“”.join(reversed(str))O(N)O(N)
Comparison of string reversal methods

Hope you learnt many Pythonic ways to reverse a string or an array of chars[]. There are many non-Pythonic ways as well which we will discuss in a separate post. Have a great day!

Posted in Arrays and Strings, Data Structures/ Leetcode, June Leetcoding Challenge

[Problem Solving] Two City Scheduling

We are given a list of costs as shown below.

costs = [[10, 20], [30, 200], [400, 50], [30, 20]]

  • The cost of flying the ith person to city A is costs[i][0]
  • The cost of flying the ith person to city B is costs[i][1]

A company is planning to interview 2N people. Return the minimum cost to fly every person to a city such that exactly N people arrive in each city.

Example:

Input: [[10, 20], [30, 200], [400, 50], [30, 20]]
Output: 110
Explanation:
The first person goes to city A for a cost of 10.
The second person goes to city A for a cost of 30.
The third person goes to city B for a cost of 50.
The fourth person goes to city B for a cost of 20.

The total minimum cost is 10 + 30 + 50 + 20 = 110 to have half the people interviewing in each city.

Assumption:

The costs list will always have an even length.

Solution:

First let’s try to understand the problem statement clearly.

  • What we need to do is out of N given people and their costs, exactly N/2 people should be flown to city A and other N/2 people need to be flown to city B.
  • This means for N/2 elements in the costs list, we have to choose the costs[i][0]th element and for the other N/2 elements, we need to choose the costs[i][1]th element.
  • The part that needs to be solved is to choose elements such that the overall cost is minimized.

The first thing that pops in our heads is to find out if we can somehow sort this given list. Let’s look at an example.

[[10, 20], [10, 20], [10, 40], [10, 50]]

We can sort by the first element of the list of lists ie., costs[i][0]. This will ensure that the list is sorted by the lowest costs to fly to city A. In this case, the list remains the same.

[[10, 20], [10, 20], [10, 40], [10, 50]]

If we choose the first N/2 people this way to fly to city A and the next N/2 people to fly to city B, we are not guaranteed that the sum of costs to fly to B will be minimum.

Sum(A) = 10 + 10 = 20

Sum(B) = 40 + 50 = 90

Total cost = 20 + 90 = 110 Wrong answer!

In order to get the optimal cost, we have to fly the first two people to city B and the next two to city A. Remember that we cannot fly the same person to both city A and city B!!

Sum(A) = 10 + 10 = 20

Sum(B) = 20 + 20 = 40

Total cost = 20 + 40 = 60 Voila!

Using the given elements in the list, we cannot directly sort it such that the cost is minimized. So, what can we do with the given values? Let’s look at the costs on a case by case basis.

[10, 20] – The cost of sending this person to city A vs city B is 10 – 20 = -10

[10, 20] – The cost of sending this person to city A vs city B is 10 – 20 = -10

[10, 40] – The cost of sending this person to city A vs city B is 10 – 20 = -30

[10, 50] – The cost of sending this person to city A vs city B is 10 – 20 = -40

Since we are subtracting cost(B) from cost(A), if the difference is low, it means it’s better to send this person to city A and when the difference is very high, we send them to city B. Now let’s say we have costs for a person as [400, 50]. The cost of sending this person to city A vs city B (cityA – cityB) is +$350!! We might as well send them to city B. So that’s what we will do.

We will sort the costs by the difference cost(A) – cost(B) from low to high. For the first N/2 people of this sorted list, we send them to city A and for the next N/2, we send them to city B.

Time Complexity: Sorting takes O(N log N), so this will be O(N log N) + O(N/2) + O (N/2). So the overall complexity is O(N log N).

Space complexity: If we are doing an in-place sort, then the space complexity will be O(1)

Code for the above solution:

def twoCitySchedCost(self, costs: List[List[int]]) -> int:

    minCost = 0
    N = len(costs)

    costs.sort(key = lambda x: x[0] - x[1])
    for i in range(N // 2):
        minCost += costs[i][0]

    for i in range(N // 2, N):    
        minCost += costs[i][1]

    return minCost

Testing:

You can test the following cases.

  • Any given test case: [[10,20],[30,200],[400,50],[30,20]]
  • Case where cost to city A is the same and vice versa: [[10,20],[10,20],[10,40],[10,50]]
  • Case where cost(A) = cost(B): [[10,10],[200,200],[400,400],[30,30]]
  • A random case mixed with the above two cases: [[20,10],[400,200],[100,400],[3,30], [3,30],[40,40]]

Hope you learnt how to solve the two city scheduling problem. Have a nice day!

Posted in Data Structures/ Leetcode, June Leetcoding Challenge, Linked Lists

[Problem Solving] Delete a linked list node when head is not given

Given a linked list as shown below, delete a node when we have access to only that node.

Input: [1,2,3,4,5] Node to delete = 4

Linked List

Output: [1,2,3,5]

Here are some of the assumptions that we will make for this problem.

  • The linked list will have at least two nodes
  • The node to delete won’t be a tail node and will be valid

Usually, when the head of a linked list is available, it is straightforward to traverse the list while keeping tracking of current and previous nodes. Once we reach the desired node to be deleted, we set its previous node’s next pointer to be the current node’s next element and re-wire some connections as shown below.

Delete a node with access to head

In this case, we don’t have access to the previous node. We have to delete it but preserve the value of the node next to it. So what we can do is to copy the value of the next node into the current node’s value and delete the next node.

Delete a node without access to head

Time Complexity: O(1)

Space Complexity: O(1)

Code for the above solution:

class ListNode:
    def __init__(self, x):
        self.val = x
        self.next = None

class Solution:
    def deleteNode(self, node):
        """
        :type node: ListNode
        :rtype: void
        """
        next_node = node.next
        node.val = next_node.val
        node.next = next_node.next
        next_node.next = None

Hope you enjoyed a simple solution to delete a linked list node when we have access only to that node. Have a great day!

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] Spiral Matrix

Given a matrix of m x n elements, return the elements in spiral order.

[
 [ 1, 2, 3 ],
 [ 4, 5, 6 ],
 [ 7, 8, 9 ]
]

Output: [1,2,3,6,9,8,7,4,5]

This problem is quite popular and is frequently asked during interviews.

Solution: To traverse the matrix in spiral order, we have to go along the boundaries starting from the top left element at i = 0, j = 0 where i, j represent the row and column indexes respectively. Let’s view the matrix as a collection of layers. Once we complete the outermost layer, we have to keep traversing the inner layers until no element is left to be visited. We can think of this as traversing multiple matrices in spiral order layer by layer.

Spiral Matrix Traversal

What are some things to remember while traversing?

  • We need some indexes to track our traversal. A straightforward approach is to have four indexes to keep track of the four boundaries of the matrix. These will be called left, right, top and bottom.
  • We will be using a while loop to add elements to a results array as long as left < right and top < bottom. This indicates that there are still elements left to traverse.
  • Inside the while loop, we will be using four for loops to traverse each boundary.
  • Avoid adding duplicate elements when changing directions ie., once we finish traversing a row [1, 2, 3, 4, 5], we have to increment the row index to avoid adding 5 again during the right most column traversal. Similarly, right and bottom pointers should be decremented upon completing their respective boundary traversals.
  • This problem code is nothing fancy but a matter of writing clear code without getting confused with the indices. The best way to get it right is to write the for loops as you traverse an example asymmetric matrix. This is because, a symmetric matrices usually do not help us handle edge cases in matrix problems.
  • One corner case for this problem is when we are traversing a column matrix as shown below where m = 4, n = 1.
Column Matrix

In this case, we first initialize these variables.

top = 0
bottom = m #4
left = 0
right = n #1

Top boundary: row = top = 0, col = left to right [0, 1). Results = [1]. Increment top pointer, top = 1.

Right boundary: row = top to bottom [1, 4), col = right – 1 = 0. Results = [ 1, 2, 3, 4]. Decrement right pointer, right = 0.

Now after the above two for loops, we are done traversing the entire matrix. But the next traversal will add duplicate elements to the results.

Bottom boundary: row = bottom – 1 = 3, col = right – 1 to left – 1 [0, -1). Results = [1, 2, 3, 4, 4]

This will attempt to traverse the bottom element 4 again. So right after traversing top and right boundaries, we have to always check whether left < right and top < bottom condition still holds true. In this case top < bottom since 1 < 4 but left = right since 0 = 0 indicating that we are about to traverse a boundary again.

Time Complexity: We will end up visiting each element of the matrix, so time complexity will be O(M * N).

Space Complexity: Constant variables are used to track boundaries, so no additional space is being used. The space complexity will be O(1).

Code:

def spiralOrder(self, matrix: List[List[int]]) -> List[int]:
        
        if not matrix or not matrix[0]:
            return []
        
        m, n = len(matrix), len(matrix[0])
        top, bottom = 0, m
        left, right = 0, n
        
        results = []
        
        while top < bottom and left < right:
            
            # traverse top boundary
            for i in range(left, right):
                results.append(matrix[top][i])
            top += 1
                
            # traverse right boundary
            for i in range(top, bottom):
                results.append(matrix[i][right-1])
            right -= 1
            
            if top < bottom and left < right:
                # traverse bottom boundary
                for i in range(right - 1, left - 1, -1):
                    results.append(matrix[bottom-1][i])
                bottom -= 1

                # traverse left boundary
                for i in range(bottom - 1, top - 1, -1):
                    results.append(matrix[i][left])
                left += 1
            
        return results

Testing: The following cases are good to test for matrix problems.

  • Empty matrix: []
  • Empty rows: [[]]
  • Single value matrix: [[1]]
  • Two element matrix: [[1, 2]]
  • Column matrix: [[1,2,3,4]]
  • Row matrix: [[1],[2],[3],[4]]
  • Symmetric matrix: [[1,2,3, 4],[4,5,6, 8], [9,10,11,12],[13,14,15,16]]
  • Asymmetric matrix: [[1,2,3, 4],[4,5,6, 8]]

Want to try solving this problem? https://leetcode.com/problems/spiral-matrix/

Hope you enjoyed this post on traversing a matrix in spiral order. Have a great day!

Posted in Arrays and Strings, Data Structures/ Leetcode, Matrix problems

[Problem Solving] Diagonal Traverse

1. Diagonal Traversal

Given a matrix of M x N elements, return the elements in diagonal order as shown below.

[
 [ 1, 2, 3, 4 ],
 [ 5, 6, 7, 8 ],
 [ 9, 10, 11, 12 ],
 [13, 14, 15, 16 ]
]

Output: [1, 2, 5, 9, 6, 3, 4, 7, 10, 13, 14, 11, 8, 12, 15, 16]

2. Simple diagonal traversal

Solution #1: To begin with, let’s ignore the part where alternate diagonal elements need to be reversed. Let’s try to traverse the matrix as shown above in image 2.

We see that a diagonal always starts either from the top row or the last column of elements. If we are given the starting point of a diagonal, how can we get to the next diagonal element?

For example, let’s say we are at 4 which is at (r = 0, c = 3) where r and c indicate the row and column indexes respectively. To get to the next element, we have to move one row down and one column to the left, ie., r = 1, c = 2. To generalize:

next_r = r + 1, next_c = c – 1

3. Traverse a diagonal

To match the expected results where alternate diagonal elements are reversed in order, we can collect the diagonal elements and reverse the order of alternate diagonals before adding to the results. The code below is an intuitive way of solving this problem.

def findDiagonalOrder(self, matrix: List[List[int]]) -> List[int]:
        
        if not matrix or not matrix[0]:
            return []
        
        m, n = len(matrix), len(matrix[0])
        results = []
        row, col = 0, 0
        
        while row < m and col < n:
            
            temp = []
            
            # To collect all elements of a diagonal, keep going one row down and column 
            # to the left until we hit matrix boundaries
            i, j = row, col
            while i < m and j >= 0:
                temp.append(matrix[i][j])
                i += 1
                j -= 1
                
            # Reverse the order for alternate diagonals
            if (row + col) % 2 == 0:
                results.extend(temp[::-1])
            else:
                results.extend(temp)
                
            # Set row, col to traverse the next diagonal
            if col < n - 1:
                col += 1
            else:
                row += 1
                
        return results

Time Complexity: O(M*N) where M is the number of rows and N is the number of columns in the matrix

Space Complexity: The additional space used in this solution is the temp array which holds the diagonal elements. The largest diagonal of a matrix has min(M, N) elements, so space complexity is O(min(M, N)).

Testing: When it comes to matrix problems, make sure to test the following cases.

  • Empty matrix: []
  • Empty rows: [[]]
  • Single value matrix: [[1]]
  • Two element matrix: [[1, 2]]
  • Column matrix: [[1,2,3,4]]
  • Row matrix: [[1],[2],[3],[4]]
  • Symmetric matrix: [[1,2,3, 4],[4,5,6, 8], [9,10,11,12],[13,14,15,16]]
  • Asymmetric matrix: [[1,2,3, 4],[4,5,6, 8]]

Solution #2: The above solution is perfectly acceptable. If you notice the indices of diagonal elements, they all sum to the same number. Remember this pattern as it might come handy while manipulating diagonal elements of a matrix.

Sum of diagonal indices

Keeping this in mind, here is an alternate solution to this problem where we use the sum of indices as keys in a dictionary to store values belonging to a particular diagonal.

import collections

def findDiagonalOrder(self, matrix: List[List[int]]) -> List[int]:
        
        if not matrix or not matrix[0]:
            return []
        
        m, n = len(matrix), len(matrix[0])
        elements = collections.defaultdict(list)
        results = []
    
        for i in range(m):
            for j in range(n):
                elements[i+j].append(matrix[i][j])
                
        for k, v in elements.items():
            if k % 2 == 0:
                results.extend(v[::-1])
            else:
                results.extend(v)

        return results

Time Complexity: O(M*N)

Space Complexity: O(M*N)

Want to try solving this problem? https://leetcode.com/problems/diagonal-traverse/

Hope you enjoyed solving the diagonal traverse problem. Have a great day!

Posted in Arrays and Strings, Data Structures/ Leetcode

[Problem Solving] Plus one to array

Given a non-empty array of digits representing a non-negative integer, plus one to the integer. The digits are stored such that the most significant digit is at the head of the list, and each element in the array contain a single digit. You may assume the integer does not contain any leading zero, except the number 0 itself.

Examples:

  • [1, 2, 3] should return [1, 2, 4]
  • [9, 9] should return [1, 0, 0]

Solution: When we add 1 to a single digit number, there are two outcomes. It can either remain a single digit number or results in 10. If the result is 10, we have to store 0 in the index i and propagate the carry 1 to the previous index i-1.

Let’s traverse the array from the end and add one to the value at index i.

  • If the value becomes 10, store 0 in nums[i] and continue traversing the array
  • Otherwise, add one to the value at nums[i] and return the array

Time Complexity: Since we traverse the array once, O(N)

Space Complexity: No additional space is used, so O(1)

def plusOne(self, digits: List[int]) -> List[int]:
        
        for i in range(len(digits)-1, -1, -1):
            
            if digits[i] + 1 == 10:
                digits[i] = 0
            else:
                digits[i] += 1
                return digits
            
        return [1] + digits #handle the case where all digits become 9

Testing: For most problems it’s good to test the following cases.

  • Empty array: []
  • Single value array: [1]
  • Array with two values: [1, 8]
  • A random array: [1, 2, 3, 4]

In this problem, there are two important cases to test. When the last digit is 9 and when all the digits are 9.

  • [1, 9]
  • [9] or [9, 9] or [9, 9, 9]

Want to try solving this problem? https://leetcode.com/problems/plus-one/

Hope you enjoyed this simple approach to the plus one problem. 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!