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

One key simplification I came up with using the hint from the problem about reducing by 1 if the number is non-zero is that we can think of any non-zero number x as x number of 1s. This means that we can convert the array into a binary-like array where each non-zero number is represented as 1 and zeros remain as 0.

From here on we can use prefix sums to calculate the number of non-zero elements to the left and right of each zero element.

For each zero element, if the count of non-zero elements on the left is equal to the count on the right, then starting from that index will lead to all elements becoming zero - so we can go in either direction and and end with a completely zeroed array.

If the counts differ by one, it means we can still start from that index but only one side will be fully reduced - the side that has one more 1.

And that’s it.

Solution in Python


class Solution:
    def countValidSelections(self, nums: List[int]) -> int:
        l = len(nums)

        prefix = [0] * l
        prefix[0] = nums[0] if nums[0] else 0

        suffix = [0] * l
        suffix[-1] = nums[-1] if nums[-1] else 0

        # note indices of any element that is `0`
        z = [i for i in range(l) if not nums[i]]

        # calculate number of `1`s to the left of this index
        for i in range(1, l):
            n = nums[i]
            v = 0

            if n:
                v += n

            prefix[i] = prefix[i - 1] + v

        # calculate number of `1`s to the right of this index
        for i in reversed(range(l - 1)):
            n = nums[i]
            v = 0

            if n:
                v += n
            
            suffix[i] = suffix[i + 1] + v

        res = 0
        
        # edge case: if all elements are `0`, we can start from any index
        if len(z) == l:
            return 2*l

        # for each zero index, check the counts on left and right
        for index in z:
            left = right = 0

            if 0 <= index - 1:
                left = prefix[index - 1]

            if index + 1 < l:
                right = suffix[index + 1] 

            # we can go either way
            if left == right:
                res += 2
            # we can go only one way
            elif left == right + 1 or right == left + 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 messed up the second case where the counts differ by one.

And we are done.