如何使用另一个ndarray的行索引ndarray？

Question

If I have

x = np.arange(1, 10).reshape((3,3))
# array([[1, 2, 3],
#        [4, 5, 6],
#        [7, 8, 9]])

and

ind = np.array([[1,1], [1,2]])
# array([[1, 1],
#        [1, 2]])

, how do I get use each row (axis 0) of ind to extract a cell of x? I hope to end up with the array [5, 6] . np.take(x, ind, axis=0) does not seem to work.

Answer 1

You could use "advanced integer indexing" by indexing x with two integer arrays, the first array for indexing the row, the second array for indexing the column:

In [58]: x[ind[:,0], ind[:,1]]
Out[58]: array([5, 6])

Answer 2

x[ind.T.tolist()]

works, too, and can also be used for multidimensional NumPy arrays.

Why?

NumPy arrays are indexed by tuples. Usually, these tuples are created implicitly by python :

Note

In Python, x[(exp1, exp2, ..., expN)] is equivalent to x[exp1, exp2, ..., expN]; the latter is just syntactic sugar for the former.

Note that this syntactic sugar isn't NumPy-specific. You could use it on dictionaries when the key is a tuple:

In [1]: d = { 'I like the number': 1, ('pi', "isn't"): 2}

In [2]: d[('pi', "isn't")]
Out[2]: 2

In [3]: d['pi', "isn't"]
Out[3]: 2

Actually, it's not even related to indexing:

In [5]: 1, 2, 3
Out[5]: (1, 2, 3)

Thus, for your NumPy array, x = np.arange(1,10).reshape((3,3))

In [11]: x[1,2]
Out[11]: 6

because

In [12]: x[(1,2)]
Out[12]: 6

So, in unutbu's answer , actually a tuple containing the columns of ind is passed:

In [21]: x[(ind[:,0], ind[:,1])]
Out[21]: array([5, 6])

with x[ind[:,0], ind[:,1]] just being an equivalent ( and recommended ) short hand notation for the same.

Here's how that tuple looks like:

In [22]: (ind[:,0], ind[:,1])
Out[22]: (array([1, 1]), array([1, 2]))

We can construct the same tuple diffently from ind : tolist() returns a NumPy array's rows. Transposing switches rows and columns, so we can get a list of columns by first transposing and calling tolist on the result:

In [23]: ind.T.tolist()
Out[23]: [[1, 1], [1, 2]]

Because ind is symmetric in your example, it is it's own transpose. Thus, for illustration, let's use

In [24]: ind_2 = np.array([[1,1], [1,2], [0, 0]])
# array([[1, 1],
#        [1, 2],
#        [0, 0]])

In [25]: ind_2.T.tolist()
Out[25]: [[1, 1, 0], [1, 2, 0]]

This can easily be converted to the tuples we want:

In [27]: tuple(ind_2.T.tolist())
Out[27]: ([1, 1, 0], [1, 2, 0])

In [28]: tuple(ind.T.tolist())
Out[28]: ([1, 1], [1, 2])

Thus,

In [29]: x[tuple(ind.T.tolist())]
Out[29]: array([5, 6])

equivalently to unutbu's answer for x.ndim == 2 and ind_2.shape[1] == 2 , but also working more generally when x.ndim == ind_2.shape[1] , in case you have to work with multi-dimensional NumPy arrays.

Why you can drop the tuple(...) and directly use the list for indexing, I don't know. Must be a NumPy thing:

In [43]: x[ind_2.T.tolist()]
Out[43]: array([5, 6, 1])