Question

是否有一种有效的方法可以对多个数组建立索引？

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

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

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

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

谢谢。

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：在多个数组上进行花式索引

2 个答案: