Julia中按特定顺序将三维数组重塑为二维数组的最优方法

问题场景 #

给定三维数组：

test_mat3 = rand(2, 3, 2)
# 输出示例：
# 2×3×2 Array{Float64, 3}:
# [:, :, 1] =
#  0.539697  0.17689   0.944783
#  0.537687  0.660573  0.19497
# 
# [:, :, 2] =
#  0.726739    0.798789  0.0900478
#  0.00744727  0.837076  0.25798

需要将其按规则重塑：把每个内层2×3矩阵的所有行拼接成1×6的数组，再将这些数组垂直拼接为2×6的二维数组，预期输出：

# 2×6 Array{Float64, 2}:
#  0.539697    0.17689   0.944783    0.537687    0.660573  0.19497
#  0.726739    0.798789  0.0900478   0.00744727  0.837076  0.25798

原实现方案（使用SplitApplyCombine包） #

已实现的代码可以正常输出预期结果：

using SplitApplyCombine

function flatten_3d(arr::Array{Float64, 3})
    mat = zeros(
        Float64,
        size(arr, 3),
        size(arr, 1)*size(arr, 2)
    )

    for (idx, arr)=enumerate(splitdimsview(arr, 3))
        the_vec = vcat(eachrow(arr)...)
        mat[idx, :] = the_vec
    end

    return mat
end

更优的内置函数实现 #

不需要依赖第三方包，用Julia内置的permutedims和reshape就能一步完成，性能也更优：

reshape(permutedims(test_mat3, (3, 1, 2)), size(test_mat3, 3), :)

分步解释：

1. permutedims(test_mat3, (3, 1, 2))：调整数组维度顺序，将原三维数组的第3维（内层矩阵的索引）移到第一位，原第1维（行）移到第二位，原第2维（列）移到第三位。调整后数组维度变为(2, 2, 3)，每个元素对应原数组中第k个内层矩阵的第i行第j列的值。
2. reshape(..., size(test_mat3, 3), :)：将调整后的三维数组重塑为二维数组，指定行数为原数组第3维的大小（2），列数由Julia自动计算（2*3=6），完美匹配需求的结构。

验证结果：

julia> reshape(permutedims(test_mat3, (3, 1, 2)), size(test_mat3, 3), :)
2×6 Matrix{Float64}:
 0.539697  0.17689   0.944783   0.537687    0.660573  0.19497
 0.726739  0.798789  0.0900478  0.00744727  0.837076  0.25798

另外也可以用eachslice结合vec实现，但性能略逊于上面的方案：

reduce(vcat, vec.(eachslice(test_mat3; dims=3))')

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