LeetCode Daily Problem P3354 Make Array Elements Equal to Zero

---

Timeline

timeline

2025-10-28

init

prefix sum

---

Problem:

P3354 Make Array Elements Equal to Zero

https://leetcode.cn/problems/make-array-elements-equal-to-zero/description/?envType=daily-question&envId=2025-10-28

According to the problem, if the sum of all elements to the left of the chosen index equals the sum of all elements to the right, then moving left or right both work. If the left sum is greater than the right sum by 1, or the right sum is greater than the left sum by 1, then only the larger side can be chosen. Use prefix sum to quickly compute the sum of array elements.

struct Solution;

impl Solution {
    pub fn count_valid_selections(nums: Vec<i32>) -> i32 {
        let mut scheme = 0;
        let mut vec: Vec<i32> = Vec::with_capacity(nums.len()); // Calculate prefix sum
        let mut last = 0;
        for &val in nums.iter() {
            last = last + val;
            vec.push(last);
        }

        let total = if let Some(&sum) = vec.last() { sum } else { 0 };

        for (index, &val) in nums.iter().enumerate() {
            if val == 0 {
                let left_sum = if let Some(&sum) = vec.get(index - 1) {
                    sum
                } else {
                    0
                };
                let right_sum = if let Some(&sum) = vec.get(index) {
                    total - sum
                } else {
                    0
                };
                if left_sum == right_sum {
                    scheme += 2;
                } else if left_sum == right_sum + 1 || right_sum == left_sum + 1 {
                    scheme += 1;
                }
            }
        }
        scheme
    }
}