跳转至

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)\)。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
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;
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
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;
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
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
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
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;
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
impl Solution {
    pub fn count_valid_selections(nums: Vec<i32>) -> i32 {
        let s: i32 = nums.iter().sum();
        let mut ans = 0;
        let mut l = 0;
        for &x in &nums {
            if x != 0 {
                l += x;
            } else if l * 2 == s {
                ans += 2;
            } else if (l * 2 - s).abs() <= 1 {
                ans += 1;
            }
        }
        ans
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
public class Solution {
    public int CountValidSelections(int[] nums) {
        int s = nums.Sum();
        int ans = 0, l = 0;
        foreach (int x in nums) {
            if (x != 0) {
                l += x;
            } else if (l * 2 == s) {
                ans += 2;
            } else if (Math.Abs(l * 2 - s) <= 1) {
                ans += 1;
            }
        }
        return ans;
    }
}

评论