使用索引列表对n维numpy数组进行切片

Question

使用索引列表对n维numpy数组进行切片

11

假设我有一个三维的numpy数组：

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

我有两个索引列表，分别对应第2和第3个维度：

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的错误，因为第一个维度的[:]似乎被忽略了。

- pbreach

3个回答

5

一种方法是使用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

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

- DSM

这就是我一直在寻找的答案！谢谢。我认为这应该被视为最普遍适用的答案，因此应该被接受。 - Clemson

2

我认为你的方法存在三个问题：

both second and third 应该是slices
由于“to”索引是排除的，它们应该从1到3和3到5
而不是A[:][second][third]，你应该使用A[:,second,third]

试试这个：

>>> np.random.seed(1145)
>>> A = np.random.random((5,5,5))                       
>>> second = slice(1,3)
>>> third = slice(3,5)
>>> A[:,second,third].shape
(5, 2, 2)
>>> A[:,second,third].flatten()
array([ 0.43285482,  0.80820122,  0.64878266,  0.62689481,  0.01298507,
        0.42112921,  0.23104051,  0.34601169,  0.24838564,  0.66162209,
        0.96115751,  0.07338851,  0.33109539,  0.55168356,  0.33925748,
        0.2353348 ,  0.91254398,  0.44692211,  0.60975602,  0.64610556])

- tobias_k

仔细检查后，我认为我误解了问题：second和third不应该是范围，而是你想要确切的那些索引--使用连续的索引有点误导人。尽管如此，我仍然会将其作为答案留在这里，以保证完整性。 - tobias_k

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- Warren Weckesser · Accepted Answer

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使用的同一块内存的视图。其中一个地方的就地更改将同时影响两个。