题目描述

给你一个整数数组 nums。

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

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

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

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

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

样例

输入：nums = [1,0,2,0,3]

输出：2

解释：

可能的有效选择方案如下：

选择 curr = 3 并向左移动。
[1,0,2,0,3] -> [1,0,2,0,3] -> [1,0,1,0,3] -> [1,0,1,0,3] -> [1,0,1,0,2] 
    -> [1,0,1,0,2] -> [1,0,0,0,2] -> [1,0,0,0,2] -> [1,0,0,0,1] -> [1,0,0,0,1]
    -> [1,0,0,0,1] -> [1,0,0,0,1] -> [0,0,0,0,1] -> [0,0,0,0,1] -> [0,0,0,0,1]
    -> [0,0,0,0,1] -> [0,0,0,0,0]
选择 curr = 3 并向右移动。
[1,0,2,0,3] -> [1,0,2,0,3] -> [1,0,2,0,2] -> [1,0,2,0,2] -> [1,0,1,0,2]
    -> [1,0,1,0,2] -> [1,0,1,0,1] -> [1,0,1,0,1] -> [1,0,0,0,1] -> [1,0,0,0,1]
    -> [1,0,0,0,0] -> [1,0,0,0,0] -> [1,0,0,0,0] -> [1,0,0,0,0] -> [0,0,0,0,0]

输入：nums = [2,3,4,0,4,1,0]

输出：0

解释：

不存在有效的选择方案。

限制

• 1 <= nums.length <= 100
• 0 <= nums[i] <= 100
• 至少存在一个元素 i 满足 nums[i] == 0。

算法

(前后缀分解) $O(n)$

1. 如果一个值为 $0$ 的位置，其前缀和后缀的和相等，则答案累加 2；如果相差为 1，则答案累加 1。
2. 预处理后缀和数组，然后遍历数组，并累加前缀和。
3. 对于值为 $0$ 的位置，按照上述条件判断。

时间复杂度

• 遍历数组两次，故时间复杂度为 $O(n)$。

空间复杂度

• 需要 $O(n)$ 的额外空间存储后缀和数组。

C++ 代码

class Solution {
public:
    int countValidSelections(vector<int>& nums) {
        const int n = nums.size();

        vector<int> suf(n + 1, 0);
        for (int i = n - 1; i >= 0; i--)
            suf[i] = suf[i + 1] + nums[i];

        int pre = 0, ans = 0;
        for (int i = 0; i < n; i++) {
            if (nums[i] > 0) {
                pre += nums[i];
            } else {
                if (pre == suf[i + 1])
                    ans += 2;
                else if (abs(pre - suf[i + 1]) == 1)
                    ++ans;
            }
        }

        return ans;
    }
};