Julia中cumsum逆运算：从多维累积和数组还原原数组

问题背景 #

给定矩阵Y是原始矩阵X经过两次累积和得到的：

Y = cumsum(cumsum(X, dims=1), dims=2)

示例如下：

julia> X = [1 4 2 3; 2 4 5 2; 4 3 4 1; 2 5 4 2];
julia> Y = cumsum(cumsum(X,dims=1), dims=2)
4x4 Matrix{Int64}:
1  5  7 10
3 11 18 23
7 18 29 35
9 25 40 48

直接使用diff函数只能还原内部元素，无法恢复第一行和第一列：

julia> diff(diff(Y, dims=1), dims=2)
3x3 Matrix{Int64}:
4 5 2
3 4 1
5 4 2

虽然可以通过拼接零矩阵实现完整还原，但这种方法会额外占用内存和时间，因此需要更高效的方案，同时要将逻辑扩展到任意维度的数组（比如三维数组Y = cumsum(cumsum(cumsum(X, dims=1), dims=2), dims=3)）。

二维数组的高效还原方案 #

无需拼接零矩阵，直接针对边界和内部元素分别处理：

1. 第一行第一列：X[1,1] = Y[1,1]
2. 第一行其余列：通过对Y的第一行做一维差分得到
3. 第一列其余行：通过对Y的第一列做一维差分得到
4. 内部元素：通过两次二维差分得到

对应Julia代码实现：

function reverse_2d_cumsum(Y)
    X = similar(Y)
    # 第一行第一列
    X[1,1] = Y[1,1]
    # 第一行
    X[1,2:end] = diff(Y[1,:])
    # 第一列
    X[2:end,1] = diff(Y[:,1])
    # 内部元素
    X[2:end,2:end] = diff(diff(Y, dims=1), dims=2)
    return X
end

测试验证：

julia> reverse_2d_cumsum(Y) == X
true

该方法仅创建与Y同大小的数组，无额外内存开销，计算效率更高。

扩展到任意维度数组 #

对于N维数组Y（由X依次在维度1到N上累积和得到），还原X的核心逻辑是：对每个元素索引(i₁,i₂,...,iₙ)，X的取值等于Y[I]交替加减所有仅减少k个索引（k从1到N）且索引仍有效（≥1）的Y值。

通用Julia实现如下：

using Combinatorics

function reverse_cumsum(Y)
    X = similar(Y)
    nd = ndims(Y)
    for I in CartesianIndices(Y)
        total = Y[I]
        idx_tuple = Tuple(I)
        # 遍历所有非空索引子集
        for k in 1:nd
            for subset in combinations(1:nd, k)
                new_idx = collect(idx_tuple)
                valid = true
                for d in subset
                    new_idx[d] -= 1
                    new_idx[d] < 1 && (valid = false; break)
                end
                valid || continue
                # 根据子集大小确定符号：k为奇数减，偶数加
                sign = isodd(k) ? -1 : 1
                total += sign * Y[CartesianIndex(new_idx...)]
            end
        end
        X[I] = total
    end
    return X
end

测试三维场景：

# 创建三维原始数组
X_3d = rand(1:5, 2,2,2)
# 生成累积和数组Y
Y_3d = cumsum(cumsum(cumsum(X_3d, dims=1), dims=2), dims=3)
# 还原并验证
reverse_cumsum(Y_3d) == X_3d

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