Python - 如何沿给定轴寻址数组？

Question

我有一个从 3D 数组读取切片的 python 脚本，如下所示:

def get_from_array(axis, start, end, array):
    if axis == 0:
        slice = array[start:end, :, :]
    elif axis == 1:
        slice = array[:, start:end, :]
    elif axis == 2:
        slice = array[:, :, start:end]
    return slice

我忍不住认为一定有更好的方法来做到这一点!有什么建议吗？

小号

Answer 1

你可以这样做:

idx = [slice(None)] * array.ndim
idx[axis] = slice(start, end)
myslice = array[tuple(idx)]