Julia中如何按指定维度动态索引AbstractMatrix实现洗牌？

要实现沿指定维度打乱数组元素且无需大量条件判断，核心是动态生成适配目标维度的索引元组，利用Julia的元组和ntuple函数可以优雅解决这个问题：

实现代码 #

using Random: shuffle

function shuffle_dim(data::AbstractArray; dims::Integer=1)
    n = size(data, dims)
    shuffled_idx = shuffle(1:n)
    # 构造索引元组：指定维度使用打乱后的索引，其余维度取全部元素
    indices = ntuple(i -> i == dims ? shuffled_idx : (:), ndims(data))
    data[indices...]
end

代码解释 #

1. 索引元组构造：ntuple会根据输入数组的维度数生成一个元组。对于每个维度位置i：
   • 如果i等于指定的dims，则使用打乱后的索引shuffled_idx
   • 否则使用(:)（Colon类型，表示选取该维度的所有元素）
2. 索引展开：indices...将元组展开为getindex的位置参数，等价于手动写data[shuffled_idx, :]（dims=1）或data[:, shuffled_idx]（dims=2），但完全动态适配任意维度。

测试示例 #

# 测试2维矩阵
mat = reshape(1:6, 2, 3)
println("原矩阵：")
display(mat)
println("\n打乱行（dims=1）：")
display(shuffle_dim(mat; dims=1))
println("\n打乱列（dims=2）：")
display(shuffle_dim(mat; dims=2))

# 测试3维数组
arr = reshape(1:24, 2, 3, 4)
println("\n打乱第三维度（dims=3）：")
display(shuffle_dim(arr; dims=3))

优势 #

• 无需任何if-else判断，代码简洁且可扩展到任意维度的AbstractArray（包括矩阵、3维数组等）
• 完全遵循Julia的索引规则，性能与手动写索引一致

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