3354. 使数组元素等于零

 

给你一个整数数组 nums 。

开始时，选择一个满足 nums[curr] == 0 的起始位置 curr ，并选择一个移动 方向 ：向左或者向右。

此后，你需要重复下面的过程：

• 如果 curr 超过范围 [0, n - 1] ，过程结束。
• 如果 nums[curr] == 0 ，沿当前方向继续移动：如果向右移，则 递增 curr ；如果向左移，则 递减 curr 。
• 如果 nums[curr] > 0:
  • 将 nums[curr] 减 1 。
  • 反转 移动方向（向左变向右，反之亦然）。
  • 沿新方向移动一步。

如果在结束整个过程后，nums 中的所有元素都变为 0 ，则认为选出的初始位置和移动方向 有效 。

返回可能的有效选择方案数目。

 

class Solution:
    def countValidSelections(self, nums: List[int]) -> int:
        left_sum = [0 for _ in nums]
        right_sum = [0 for _ in nums]

        left_sum[0] = 0
        right_sum[0] = 0

        for i,n in enumerate(nums):
            if i>0:
                left_sum[i] = left_sum[i-1] + nums[i-1]
        for i,n in enumerate(nums[::-1]):
            if i>0:
                right_sum[i] = right_sum[i-1] + nums[::-1][i-1]
        right_sum = right_sum[::-1]
        # print(left_sum)
        # print(right_sum)
        ans = 0
        for i,n in enumerate(nums):
            # print(i,n)
            if n==0 and left_sum[i]==right_sum[i]:
                ans +=  2
            elif n==0 and abs(left_sum[i]-right_sum[i])==1:
                ans += 1
        return ans