3028. 边界上的蚂蚁 Ant on the Boundary

问题描述

这是一个经典的前缀和数组问题.

解题思路如下:

1. 先遍历 nums 数组, 构造出前缀和数组 positions
2. 遍历 positions 数组, 统计里面数值 0 出现的次数

该算法的实现如下:

#![allow(unused)]
fn main() {
// Prefix Sum
pub fn return_to_boundary_count1(nums: Vec<i32>) -> i32 {
    debug_assert!(!nums.is_empty());

    let mut positions = vec![0; nums.len()];
    positions[0] = nums[0];
    for i in 1..nums.len() {
        positions[i] = positions[i - 1] + nums[i];
    }
    positions.into_iter().filter(|&x| x == 0).count() as i32
}
}

该算法的时间复杂度是 O(n), 空间复杂度是 O(n).

上面的算法可以做一些简化, 因为是一次性获取结果, 其是是不需要将中间值存入到 positions 数组的, 在遍历 nums 数组时可以直接更新 zeros_count 计数值. 算法实现如下:

#![allow(unused)]
fn main() {
// Prefix Sum
// 不保存中间值到数组
pub fn return_to_boundary_count2(nums: Vec<i32>) -> i32 {
    debug_assert!(!nums.is_empty());

    let mut zeros_count = 0;
    let mut last_position = 0;
    for num in nums {
        last_position += num;
        if last_position == 0 {
            zeros_count += 1;
        }
    }
    zeros_count
}
}

该算法的时间复杂度是 O(n), 空间复杂度是 O(1).