3354. 使数组元素等于零

题目描述

给你一个整数数组 nums 。

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

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

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

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

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

示例 1：

输入：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].

示例 2：

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

输出：0

解释：

不存在有效的选择方案。

提示：

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

解法

方法一：枚举 + 前缀和

假设我们初始向左移动，遇到了一个非零元素，那么我们就需要将这个元素减一，然后改变移动方向，继续移动。

因此，我们可以维护每个零值元素左侧的元素和 $l$，右侧元素的和 $s - l$。如果 $l = s - l$，即左侧元素和等于右侧元素和，那么我们可以选择当前零值元素，向左或向右移动，答案加 $2$；如果 $|l - (s - l)| = 1$，此时如果左侧元素和更大，那么我们可以选择当前零值元素，向左移动，答案加 $1$，如果右侧元素和更大，那么我们可以选择当前零值元素，向右移动，答案加 $1$。

时间复杂度 $O(n)$，其中 $n$ 为数组长度。空间复杂度 $O(1)$。

Python3JavaC++GoTypeScript

class Solution:
    def countValidSelections(self, nums: List[int]) -> int:
        s = sum(nums)
        ans = l = 0
        for x in nums:
            if x:
                l += x
            elif l * 2 == s:
                ans += 2
            elif abs(l * 2 - s) == 1:
                ans += 1
        return ans

class Solution {
    public int countValidSelections(int[] nums) {
        int s = Arrays.stream(nums).sum();
        int ans = 0, l = 0;
        for (int x : nums) {
            if (x != 0) {
                l += x;
            } else if (l * 2 == s) {
                ans += 2;
            } else if (Math.abs(l * 2 - s) <= 1) {
                ++ans;
            }
        }
        return ans;
    }
}

class Solution {
public:
    int countValidSelections(vector<int>& nums) {
        int s = accumulate(nums.begin(), nums.end(), 0);
        int ans = 0, l = 0;
        for (int x : nums) {
            if (x) {
                l += x;
            } else if (l * 2 == s) {
                ans += 2;
            } else if (abs(l * 2 - s) <= 1) {
                ++ans;
            }
        }
        return ans;
    }
};

func countValidSelections(nums []int) (ans int) {
    l, s := 0, 0
    for _, x := range nums {
        s += x
    }
    for _, x := range nums {
        if x != 0 {
            l += x
        } else if l*2 == s {
            ans += 2
        } else if abs(l*2-s) <= 1 {
            ans++
        }
    }
    return
}

func abs(x int) int {
    if x < 0 {
        return -x
    }
    return x
}

function countValidSelections(nums: number[]): number {
    const s = nums.reduce((acc, x) => acc + x, 0);
    let [ans, l] = [0, 0];
    for (const x of nums) {
        if (x) {
            l += x;
        } else if (l * 2 === s) {
            ans += 2;
        } else if (Math.abs(l * 2 - s) <= 1) {
            ++ans;
        }
    }
    return ans;
}

3354. 使数组元素等于零

题目描述

解法

方法一：枚举 + 前缀和

评论