python - 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

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])]

这在 a (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]])