Problem Statement in English

You’re given a 0-indexed integer array nums. You need to choose a starting index such that the value at that index is 0.

From that index you can either move right or left.

If the new index is within the bounds of the array, you can perform the following operation:

  • If the value at the new index is greater than 0, decrease it by 1.
  • If the value at the new index is 0, you must move to the next index in the direction you were going.

You need to continue this process until you move out of the bounds of the array.

You need to return the number of indices that you can start at such that at the end of the process, all elements in the array become 0.

Approach

It really helps if you think about it like ping pong. For every reduction you make, you’ll need to switch directions. So you’re going to need more or less the same number of pings as pongs…not entirely accurate, and you’ll see why.

There are 2 cases that you should see in this.

The first is that you have something like [5, 0, 5]. Here you have 2 ways to make the entire array 0. You can go either left or right - it doesn’t matter.

The next is something lile [5, 0, 4]. Here you have to go left first, or it’s not going to work out.

Now given this, we know that we need to be able to efficiently look up the number of pings (sum of elements on the left), and pongs (sum of elements on the right). Prefix and suffix sums to the rescue!

There’s a bunch of edge cases that need to be handled.

  • What if there’s only 1 element? The questions says it’s guaranteed that there’s atleast one $0$. Then the answer is $2$, since you can go either way.
  • What if there’sa bunch of $0$s and just one $1$? Then your answer is the number of $0$s.
  • What if it’s just $0$s? Then the answer is $2 * N$.

And we’re done!

Solution in Python


class Solution:
    def countValidSelections(self, nums: List[int]) -> int:
        N = len(nums)
        s = sum(nums)
        
        res = 0

        if N == 1:
            return 2
        elif N > 1 and s == 1:
            return N - 1
        elif not s:
            return 2 * N

        prefix = [0] * N
        prefix[0] = nums[0]

        suffix = [0] * N
        suffix[-1] = nums[-1]

        for i in range(1, N - 1):
            prefix[i] = prefix[i - 1] + nums[i]
        
        for i in reversed(range(1, N - 1)):
            suffix[i] = suffix[i + 1] + nums[i]

        for i in range(1, N - 1):
            if nums[i] == 0:
                if prefix[i - 1] == suffix[i + 1]:
                    res += 2      
                elif abs(prefix[i - 1] - suffix[i + 1]) == 1:
                    res += 1

        return res

Complexity

  • Time: $O(n)$
    Since we traverse the array a constant number of times.

  • Space: $O(n)$
    Since we build a prefix sum

Mistakes I Made

I missed literally every edge case 💀

And we are done.