如何在Julia中高效生成累加至/从(i,j)的矩阵及张量扩展

一、二维数组的更优实现 #

你的现有代码中，g(P) 使用了两次 reverse! 原地反转原数组，这会修改输入的P，导致后续使用P时得到的是反转后的值，这通常不是预期行为。可以改用非原地的reverse函数，避免破坏原数组：

f(P) = cumsum(cumsum(P, dims=1), dims=2)
g(P) = reverse(reverse(f(reverse(reverse(P, dims=1), dims=2)), dims=1), dims=2)

如果追求代码简洁性，还可以用rotr90替代两次维度反转，效果完全一致：

g(P) = f(rotr90(P, 2)) |> x -> rotr90(x, 2)

Julia内置的cumsum已经是高度优化的底层实现，二维场景下很难有比它更快的方案，上述优化主要是避免原地修改输入数组的副作用。

二、d维张量的扩展实现 #

对于任意d维张量，我们可以通过遍历所有维度依次做累积和来生成A；对于B，则先反转每个维度，调用A的生成逻辑，再反转回原维度顺序即可：

生成A（从原点到当前索引的累积和） #

function cumulative_sum_all_dims(P)
    A = copy(P)
    for dim in 1:ndims(A)
        cumsum!(A, A, dims=dim)
    end
    return A
end

使用cumsum!原地修改数组，减少内存分配，提升大张量场景下的性能。

生成B（从当前索引到末尾的累积和） #

function reverse_cumulative_sum_all_dims(P)
    reversed_P = reverse(P; dims=1:ndims(P))
    reversed_A = cumulative_sum_all_dims(reversed_P)
    return reverse(reversed_A; dims=1:ndims(P))
end

3维张量示例验证 #

# 生成3×2×2的随机张量
P = rand(1:6, 3, 2, 2)
A = cumulative_sum_all_dims(P)
B = reverse_cumulative_sum_all_dims(P)

该实现适配任意维度的张量，完全复用Julia内置的高效操作，保证性能与通用性。

内容的提问来源于stack exchange，提问作者Sakurai.JJ