CODEWITHSUNDEEP
numpy
James Arten
James Arten

Reputation: 666

Combining two `numpy` arrays by using one as index for columns

If I have two numpy arrays

arr1
Out  [7]: array([1, 0, 1, ..., 1, 0, 0])

and

arr2
Out  [6]: 
array([[0.10420547, 0.8957946 ],
       [0.6609819 , 0.3390181 ],
       [0.16680466, 0.8331954 ],
       ...,
       [0.27138624, 0.7286138 ],
       [0.6883444 , 0.31165552],
       [0.70164204, 0.298358  ]], dtype=float32)

what is the quickest way to return a new array arr3 in such a way that arr1 indicates the column that I want from arr2 for each row? I would like to return something like:

arr3
array([0.8957946, 0.6609819, 0.8331954, ... ])

I would do it by filling a new empty array and iterating but I can't think of a quicker way right now.

EDIT:

Ok, a way that I found is the following, but probably not optimal (?):

arr3 = np.array([arr2[i][arr1[i]] for i in range(len(arr2))])

returns

arr3
Out  [23]: 
array([0.8957946 , 0.6609819 , 0.8331954 , ..., 0.7286138 , 0.6883444 ,
       0.70164204], dtype=float32)

Upvotes: 0

Views: 37

Answers (1)

MSS
MSS

Reputation: 3623

You can do it like this:

np.take_along_axis(arr2,arr1[:,None],1).squeeze()

Upvotes: 1

Related Questions