Julia中实现遇0重置的累积求和（cumsum）方案问询

我们需要实现按列累积求和且遇0时重置的功能：当某元素为0时，该位置的求和结果为0，后续求和从0重新开始；非0元素则延续前一个非重置的求和结果累加。

示例输入 #

a = [1.0 2.0 3.0; 0.0 5.0 6.0; 10.0 0.0 9.0]

期望输出 #

3×3 Matrix{Float64}:
 1.0  2.0   3.0
 0.0  7.0   9.0
10.0  0.0  18.0

Julia内置的cumsum(a, dims=1)无法实现该逻辑，它会持续累加所有元素，不会在遇到0时重置：

julia> cumsum(a, dims=1)
3×3 Matrix{Float64}:
  1.0  2.0   3.0
  1.0  7.0   9.0
 11.0  7.0  18.0

解决方案 #

1. 基础实现（逐列处理）

定义一个针对一维数组的重置求和函数，再通过mapslices应用到矩阵的每一列：

function cumsum_reset0(x::AbstractVector)
    res = similar(x)
    current_sum = zero(eltype(x))
    for i in eachindex(x)
        # 若处理浮点数，建议用isapprox判断精度，避免浮点误差影响判断
        if isapprox(x[i], zero(eltype(x)); atol=1e-10)
            current_sum = zero(eltype(x))
            res[i] = zero(eltype(x))
        else
            current_sum += x[i]
            res[i] = current_sum
        end
    end
    return res
end

# 应用到矩阵的每一列
result = mapslices(cumsum_reset0, a, dims=1)

2. 高效实现（直接操作视图）

对于大型矩阵，mapslices会有额外开销，直接通过视图操作每一列能显著提升性能：

function cumsum_reset0!(res::AbstractVector, x::AbstractVector)
    current_sum = zero(eltype(x))
    for i in eachindex(x, res)
        if isapprox(x[i], zero(eltype(x)); atol=1e-10)
            current_sum = zero(eltype(x))
            res[i] = zero(eltype(x))
        else
            current_sum += x[i]
            res[i] = current_sum
        end
    end
    return res
end

# 初始化结果矩阵，逐列处理
result = similar(a)
for j in 1:size(a, 2)
    cumsum_reset0!(view(result, :, j), view(a, :, j))
end

内容的提问来源于stack exchange，提问作者Thomas Jalabert