Numpy：对多个数组进行花式索引

Question

是否有一种有效的方法来索引多个数组？

例如，我有一个要从中索引的数组

a = [[1,2,3],[4,5,6]]

另一个数组包含索引。 b = [[0, 1], [1,2]]

我期望 [[1, 2], [5, 6]]，它通过 [0,1] 对 a 的第一行进行索引，并通过 [1,2] 对 a 的第二行进行索引。

谢谢。

Answer 1

如果 a 和 b 长度相同，您可以尝试使用 np.take 如下：

import numpy as np

a = [[1,2,3],[4,5,6]]
b = [[0, 1], [1,2]]
result = [np.take(a[i],b[i]).tolist() for i in range(len(a))]

print(result)
# result: [[1, 2], [5, 6]]

Answer 2

In [107]: a = [[1,2,3],[4,5,6]]
In [108]: b = [[0, 1], [1,2]]

a 和 b 是列表。合适的解决方案是嵌套列表理解

In [111]: [[a[i][j] for j in x] for i,x in enumerate(b)]
Out[111]: [[1, 2], [5, 6]]

现在如果a做成一个numpy数组：

In [112]: np.array(a)[np.arange(2)[:,None], b]
Out[112]: 
array([[1, 2],
       [5, 6]])

为此，数组的第一个维度使用 (2,1) 数组进行索引，第二个维度使用 (2,2) 进行索引。他们一起广播产生一个 (2,2) 结果。

Numpy extract submatrix

正朝着相同的方向努力，但接受的答案使用 ix_

Y[np.ix_([0,3],[0,3])]

在 (2,2) b 的情况下不起作用。

In [113]: np.array(a)[np.ix_(np.arange(2), b)]
ValueError: Cross index must be 1 dimensional

ix_ 会将第一个维度 np.arange(2) 向右转 (2,1)。

这可能会使广播更加明确：

In [114]: np.array(a)[[[0,0],[1,1]], [[0,1],[1,2]]]
Out[114]: 
array([[1, 2],
       [5, 6]])

它选择元素 (0,0)、(0,1)、(1,1) 和 (1,2)

为了进一步测试，使 b 非对称：

In [138]: b = [[0, 1,1], [1,2,0]]       # (2,3)
In [139]: np.array(a)[np.arange(2)[:,None], b]
Out[139]: 
array([[1, 2, 2],
       [5, 6, 4]])

Numpy：对多个数组进行花式索引

Numpy: Fancy Indexing over Multiple arrays

python

numpy

vectorization

scientific-computing