2025年10月28日 星期二

accumulate/sum 加總

今天練習的題目是 3354. Make Array Elements Equal to Zero,題目可以利用加總來解題。

首先是 C++ 要使用 accumulate

class Solution {
public:
    int countValidSelections(vector< int>& nums) {
        int ans = 0, prev = 0, post = accumulate(nums.begin(), nums.end(), 0);
        for (auto num : nums) {
            if (num == 0) {
                if (prev == post)
                    ans += 2;
                else if (abs(prev - post) == 1)
                    ++ans;
            } else {
                post -= num;
                prev += num;
            }
        }
        return ans;
    }
};

Python3 就用 sum()

class Solution:
    def countValidSelections(self, nums: List[int]) -> int:
        ans, prev, post = 0, 0, sum(nums)
        for num in nums:
            if num == 0:
                if prev == post:
                    ans += 2
                elif abs(prev - post) == 1:
                    ans += 1
            else:
                prev += num
                post -= num
        return ans

Go 加總要自己來,abs() 也要自己刻,不過發現可以利用 slices.Values() 在 range 上面,這樣就不用多寫一個 _ 底線,像是 for _, num := range nums 這樣,不過多寫 slices.Values 也不一定好到哪裡去。

func abs(num int) int {
    if num < 0 {
        return -num
    }
    return num
}

func countValidSelections(nums []int) int {
    ans, prev, post := 0, 0, 0
    for num := range slices.Values(nums) {
        post += num
    }
    for num := range slices.Values(nums) {
        if num == 0 {
            if prev == post {
                ans += 2
            } else if abs(prev-post) == 1 {
                ans += 1
            }
        } else {
            prev += num
            post -= num
        }
    }
    return ans
}

最後 Rust 也可以用 sum() 不過要用在 iter() 後面,另外要注意的一點是 sum() 會需要指定型別,不然會有型別無法推導的問題,也可以用 nums.iter().sum::< i32>(),不過也可以用 fold 來算,像是 nums.iter().fold(0, |x, y| x + y) 這樣,就能成功推導型別。

nums.iter().sum() 的問題在於它非常泛型。sum() 方法可以將 &i32 的迭代器加總成多種不同的型別,例如 i32, i64, i128 等(任何實作了 Sum<&i32> trait 的型別)。

impl Solution {
    pub fn count_valid_selections(nums: Vec<i32>) -> i32 {
        let (mut ans, mut prev, mut post): (i32, i32, i32) = (0, 0, nums.iter().sum());
        for num in nums {
            if num == 0 {
                if prev == post {
                    ans += 2;
                } else if (prev-post).abs() == 1 {
                    ans += 1;
                }
            } else {
                prev += num;
                post -= num;
            }
        }
        ans
    }
}

沒有留言: