python - Numpy 2d array - take N elements from specified index

Question

Assuming I have a 2d numpy array:

mat=[[5,5,3,6,3],
     [3,2,7,8,1],
     [7,5,5,2,0]]

and a vector of indexes:

vec=[3,1,2]

What I need is to take 3 elements from the corresponding index. For example the first element in the vector, that corresponds to the first line in the matrix is 3. Thus I need to take 3 elements from index 3 (0-based) in the first line, which is 6. So what I need is [6,3,None].

The final output should be:

[[6,3,None],
 [2,7,8],
 [5,2,0]]

I tried to play with take and with fancy indexing, but couldn't get the desired output.

Any help would be appreciated!

question from:https://stackoverflow.com/questions/65921824/numpy-2d-array-take-n-elements-from-specified-index

Answer 1

You can do this -

import numpy as np

mat=np.array([[5,5,3,6,3],
            [3,2,7,8,1],
            [7,5,5,2,0]])

mat = np.hstack((mat, np.ones((3,3))*np.nan))

vec=np.array([3,1,2])
idx = vec[:, None] + np.arange(0, 3)
print(mat[np.arange(3)[:,None], idx])

Gives -

[[ 6.  3. nan]
 [ 2.  7.  8.]
 [ 5.  2.  0.]]

First just append the original array with three columns of inf or Noneor something. Then create a 2d index array from the vec by adding sequential integers from 0 and simply index the original matrix.