获取numpy数组中第k维的第i个片

Question

我有一个n-dimensional numpy数组，我想获得k-th维度的i-th切片。一定有比这更好的东西

# ... 
elif k == 5:
    b = a[:, :, :, :, :, i, ...]
# ...

Answer 1

b = a[(slice(None),) * k + (i,)]

手动构建索引元组。

如Python language reference中所述，以下形式的表达式

a[:, :, :, :, :, i]

转换为

a[(slice(None), slice(None), slice(None), slice(None), slice(None), i)]

我们可以通过直接构建元组来达到同样的效果，而不是使用切片表示法。(有一个小问题，即构建元组直接生成a[(i,)]，而不是k=0的a[i]，但NumPy对标量i的处理方式与此相同。)

Answer 2

我不确定它是否适用于k- dim，但它适用于2- dim

a.take(i,axis=k)

Answer 3

基本上，您希望能够以编程方式创建元组:, :, :, :, :, i, ...，以便将其作为a的索引传递。不幸的是，您不能简单地对冒号操作符直接使用普通的元组乘法(即，(:,) * k不能生成k冒号操作符的元组)。但是，您可以使用colon = slice(None)获取“冒号切片”的实例。然后，您可以执行b = a[(colon,) * k + (i,)]，这将有效地在ith维度的kth列对a进行索引。

将其封装在一个函数中，您将得到：

def nDimSlice(a, k, i):
    colon = slice(None)
    return a[(colon,) * k + (i,)]