使用索引列表切割n维numpy数组

Question

说我有一个3维numpy数组:

np.random.seed(1145)
A = np.random.random((5,5,5))

我有两个与第二和第三维相对应的索引列表:

second = [1,2]
third = [3,4]

我想选择numpy数组中对应的元素

A[:][second][third]

所以切片阵列的形状将是(5,2,2)和

A[:][second][third].flatten()

相当于:

In [226]:

for i in range(5):
    for j in second:
        for k in third:
            print A[i][j][k]

0.556091074129
0.622016249651
0.622530505868
0.914954716368
0.729005532319
0.253214472335
0.892869371179
0.98279375528
0.814240066639
0.986060321906
0.829987410941
0.776715489939
0.404772469431
0.204696635072
0.190891168574
0.869554447412
0.364076117846
0.04760811817
0.440210532601
0.981601369658

有没有办法以这种方式切割numpy数组？到目前为止,当我尝试时,A[:][second][third]我得到IndexError: index 3 is out of bounds for axis 0 with size 2因为[:]第一个维度似乎被忽略了.

Answer 1

Numpy使用多个索引,因此A[1][2][3]您可以 - 而且应该 - 使用A[1,2,3].

你可能认为你可以做A[:, second, third],但是numpy索引是广播的,广播second和third(两个一维序列)最终是numpy等价的zip,所以结果已经形成(5, 2).

你真正想要的是实际上second和和的外部产品一起索引third.您可以通过制作其中一个来进行广播,比如说second成形为(2,1)的二维数组.那么广播second和third一起产生的形状就是(2,2).

例如:

In [8]: import numpy as np

In [9]: a = np.arange(125).reshape(5,5,5)

In [10]: second = [1,2]

In [11]: third = [3,4]

In [12]: s = a[:, np.array(second).reshape(-1,1), third]

In [13]: s.shape
Out[13]: (5, 2, 2)

需要注意的是,在这个具体的例子,在价值观second和third是连续的.如果这是典型的,您可以简单地使用切片:

In [14]: s2 = a[:, 1:3, 3:5]

In [15]: s2.shape
Out[15]: (5, 2, 2)

In [16]: np.all(s == s2)
Out[16]: True

这两种方法有一些非常重要的区别.

• 第一种方法也适用于与切片不等效的索引.例如,它会起作用second = [0, 2, 3].(有时你会看到这种索引方式被称为"花式索引".)
• 在第一种方法(使用广播和"花式索引")中,数据是原始数组的副本.在第二种方法(仅使用切片)中,数组s2是对使用的同一内存块的视图a.一个就地改变将改变它们.

Answer 2

一种方法是使用np.ix_：

>>> out = A[np.ix_(range(A.shape[0]),second, third)]
>>> out.shape
(5, 2, 2)
>>> manual = [A[i,j,k] for i in range(5) for j in second for k in third]
>>> (out.ravel() == manual).all()
True

缺点是您必须明确指定缺失的坐标范围，但您可以将其包装到一个函数中。