给定索引列表，将任意维度的 numpy ndarry 切片为 1d 数组答案

【问题标题】：Slice numpy ndarry of arbitrary dimension to 1d array given a list of indices给定索引列表，将任意维度的 numpy ndarry 切片为 1d 数组
【发布时间】：2023-03-23 04:59:01
【问题描述】：

我有一个 numpy ndarray arr 和 indices，一个指定特定条目的索引列表。具体来说：

arr = np.arange(2*3*4).reshape((2,3,4))
indices= [1,0,3]

我有代码通过arr 获取一维切片，观察除一个索引之外的所有索引n：

arr[:, indices[1], indices[2]]  # n = 0
arr[indices[0], :, indices[2]]  # n = 1
arr[indices[0], indices[1], :]  # n = 2

我想更改我的代码以循环 n 并支持任意维度的 arr。

我查看了文档中的indexing routines 条目，并找到了有关slice() 和np.s_() 的信息。我能够拼凑出我想要的东西：

def make_custom_slice(n, indices):
    s = list()
    for i, idx in enumerate(indices):
        if i == n:
            s.append(slice(None))
        else:
            s.append(slice(idx, idx+1))
    return tuple(s)


for n in range(arr.ndim):
    np.squeeze(arr[make_custom_slice(n, indices)])

np.squeeze 用于删除长度为 1 的轴。没有这个，这个生成的数组的形状是 (arr.shape[n],1,1,...) 而不是 (arr.shape[n],)。

有没有更惯用的方法来完成这项任务？

【问题讨论】：

“有没有更惯用的方法来完成这项任务？”使用 Numpy 完成任务的惯用方法通常涉及不编写循环。
我同意这一点。我在下面发布的改进删除了构建切片对象的函数中的循环。我不确定是否可以在 n 中对函数进行矢量化以删除调用它的循环。

标签： python numpy slice numpy-ndarray numpy-slicing

【解决方案1】：

对上述解决方案的一些改进（可能仍然存在单行或更高性能的解决方案）：

def make_custom_slice(n, indices):
    s = indices.copy()
    s[n] = slice(None)
    return tuple(s)


for n in range(arr.ndim):
    print(arr[make_custom_slice(n, indices)])

整数值idx 可用于替换切片对象slice(idx, idx+1)。因为大多数索引都是直接复制的，所以从索引的副本开始，而不是从头开始构建列表。

以这种方式构建时，arr[make_custom_slice(n, indices) 的结果具有预期的维度，而np.squeeze 是不必要的。

【讨论】：