Julia中实现类Numpy通用数组切片以保留维度的方法问询

我完全懂你的痛点！Julia默认的索引行为和Numpy确实不一样，处理高维数组时手动写一堆:,:,:简直繁琐到爆炸。不过别担心，有几种简洁的方法能实现你想要的通用维度保留切片，尤其是内置函数selectdim，完美对标Numpy的用法！

方法1：使用内置的selectdim函数（首推） #

Julia专门提供了selectdim(arr, dim, idx)函数，用来在指定维度上选取索引，并且自动保留剩余维度的结构，完全符合你想要的简洁性和通用性。

举几个例子看看效果：

# 3维数组测试
arr3 = ones(3, 4, 5)
size(selectdim(arr3, 1, 3))  # 输出 (4, 5)，和Numpy里arr[2].shape对应（注意Julia索引从1开始）

# 5维数组测试
arr5 = ones(3, 4, 5, 2, 2)
size(selectdim(arr5, 1, 2))  # 输出 (4, 5, 2, 2)

# 还能灵活切其他维度，比如取第二维的第2个元素
size(selectdim(arr3, 2, 2))  # 输出 (3, 5)

方法2：用ntuple动态生成通用索引 #

如果你不想用内置函数，也可以通过ntuple自动生成对应数量的冒号，实现通用切片逻辑：

# 封装一个函数，取第一维的第k个元素并保留维度
function slice_first_dim(arr, k)
    return arr[k, ntuple(_ -> Colon(), ndims(arr)-1)...]
end

# 测试一下
size(slice_first_dim(arr3, 3))  # 输出 (4, 5)
size(slice_first_dim(arr5, 2))  # 输出 (4, 5, 2, 2)

这里ntuple(_ -> Colon(), ndims(arr)-1)会根据数组的维度数，自动生成刚好足够的冒号，再用...展开作为后续索引，不管数组是3维还是5维都能自动适配。

方法3：自定义通用切片函数（支持任意维度） #

如果需要更灵活地选择要切片的维度，可以封装一个通用函数：

function slice_keep_dims(arr, dim::Int, idx)
    # 生成索引元组：指定维度用传入的idx，其余维度全用冒号
    indices = ntuple(i -> i == dim ? idx : Colon(), ndims(arr))
    return arr[indices...]
end

# 使用示例
size(slice_keep_dims(arr3, 1, 3))  # 输出 (4, 5)
size(slice_keep_dims(arr5, 3, 4))  # 取第三维的第4个元素，输出 (3, 4, 2, 2)

总的来说，selectdim是最简洁、最原生的解决方案，完全能替代Numpy里arr[k]的用法，再也不用为高维数组写一堆冒号了！

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