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并向左移动。<ul> <li><code>[1,0,2,<strong><u>0</u></strong>,3] -> [1,0,<strong><u>2</u></strong>,0,3] -> [1,0,1,<strong><u>0</u></strong>,3] -> [1,0,1,0,<strong><u>3</u></strong>] -> [1,0,1,<strong><u>0</u></strong>,2] -> [1,0,<strong><u>1</u></strong>,0,2] -> [1,0,0,<strong><u>0</u></strong>,2] -> [1,0,0,0,<strong><u>2</u></strong>] -> [1,0,0,<strong><u>0</u></strong>,1] -> [1,0,<strong><u>0</u></strong>,0,1] -> [1,<strong><u>0</u></strong>,0,0,1] -> [<strong><u>1</u></strong>,0,0,0,1] -> [0,<strong><u>0</u></strong>,0,0,1] -> [0,0,<strong><u>0</u></strong>,0,1] -> [0,0,0,<strong><u>0</u></strong>,1] -> [0,0,0,0,<strong><u>1</u></strong>] -> [0,0,0,0,0]</code>.</li> </ul> </li> <li>选择 <code>curr = 3</code> 并向右移动。 <ul> <li><code>[1,0,2,<strong><u>0</u></strong>,3] -> [1,0,2,0,<strong><u>3</u></strong>] -> [1,0,2,<strong><u>0</u></strong>,2] -> [1,0,<strong><u>2</u></strong>,0,2] -> [1,0,1,<strong><u>0</u></strong>,2] -> [1,0,1,0,<strong><u>2</u></strong>] -> [1,0,1,<strong><u>0</u></strong>,1] -> [1,0,<strong><u>1</u></strong>,0,1] -> [1,0,0,<strong><u>0</u></strong>,1] -> [1,0,0,0,<strong><u>1</u></strong>] -> [1,0,0,<strong><u>0</u></strong>,0] -> [1,0,<strong><u>0</u></strong>,0,0] -> [1,<strong><u>0</u></strong>,0,0,0] -> [<strong><u>1</u></strong>,0,0,0,0] -> [0,0,0,0,0].</code></li> </ul> </li>
示例 2：
输入：nums = [2,3,4,0,4,1,0]
输出：0
解释：
不存在有效的选择方案。
提示：
1 <= nums.length <= 1000 <= 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)$。
Python3
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
Java
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;
}
}
C++
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;
}
};
Go
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
}
TypeScript
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;
}