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 side is greater than the right side by 1, or the right side is greater than the left side by 1, then you can only choose the larger side. Use prefix sums 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()); // Compute prefix sums
        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
    }
}