给定每一行对应的索引，从矩阵中获取对应的元素

Question

Given indexes for each row, how to return the corresponding elements in a 2-d matrix? For instance, In array of np.array([[1,2,3,4],[4,5,6,7]]) I expect to see the output [[1,2],[4,5]] given indxs = np.array([[0,1],[0,1]]) . Below is what I've tried:

a= np.array([[1,2,3,4],[4,5,6,7]])
indxs = np.array([[0,1],[0,1]]) #means return the elements located at 0 and 1 for each row
#I tried this, but it returns an array with shape (2, 2, 4) 
a[idxs]

Answer 1

The reason you are getting two times your array is that when you do a[[0,1]] you are selecting the rows 0 and 1 from your array a, which are indeed your entire array.

In[]: a[[0,1]]

Out[]: array([[1, 2, 3, 4],
       [4, 5, 6, 7]])

You can get the desired output using slides. That would be the easiest way.

a = np.array([[1,2,3,4],[4,5,6,7]])
a[:,0:2]

Out []: array([[1, 2],
               [4, 5]])

In case you are still interested on indexing, you could also get your output doing:

In[]: [list(a[[0],[0,1]]),list(a[[1],[0,1]])]

Out[]: [[1, 2], [4, 5]]

The NumPy documentation gives you a really nice overview on how indexes work.

Answer 2

In [120]: indxs = np.array([[0,1],[0,1]])
In [121]: a= np.array([[1,2,3,4],[4,5,6,7]])
     ...: indxs = np.array([[0,1],[0,1]]) #

You need to provide an index for the first dimension, one that broadcasts with with indxs .

In [122]: a[np.arange(2)[:,None], indxs]
Out[122]: 
array([[1, 2],
       [4, 5]])

indxs is (2,n), so you need a (2,1) array to give a (2,n) result