CODEWITHSUNDEEP
pythonarraysnumpy
user44136
user44136

Reputation: 25

Numpy matrix: select column based on list

I have a numpy array and a list that list that defines the rows I want to select. What is the best way to do this operation?

import numpy as np

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

b = np.array([[1],
              [0],
              [2]])

Desired result

np.array([[2],
         [4],
         [9]])

I have tried np.take() but this does not work.

Kind regards

EDIT: as this needs to be done repeatedly on a large array, I'm looking for a vectorized approach (without loops)

Upvotes: 1

Views: 487

Answers (4)

hpaulj
hpaulj

Reputation: 231355

Recent versions have added a take_along_axis which does what you want:

In [96]: a = np.array([[1,2,3], 
    ...:               [4,5,6], 
    ...:               [7,8,9]]) 
    ...:  
    ...: b = np.array([[1], 
    ...:               [0], 
    ...:               [2]])                                                                           
In [97]: np.take_along_axis(a,b,axis=1)                                                                
Out[97]: 
array([[2],
       [4],
       [9]])

It works much like @Nils answer, a[np.arange(3), np.squeeze(b)], but handles the dimensions better.

Similar recent questions:

efficient per column matrix indexing in numpy

keep elements of an np.ndarray by values of another np.array (vectorized)

Upvotes: 1

Nils Werner
Nils Werner

Reputation: 36739

If you remove the extraneous dimensions from b

b = np.sqeeze(b)

You can use the following:

a[np.arange(len(b)), b]

Upvotes: 2

sltzgs
sltzgs

Reputation: 166

It's not very pythonic but this should do the trick for your problem:

res = np.zeros(len(b))
for i, row in enumerate(a):
    res[i] = row[b[i]]

print(res)

same in one line:

a[[i[0] for i in b],[i for i in range(len(b))]]

Upvotes: 1

Serge Ballesta
Serge Ballesta

Reputation: 148880

You could use a list comprehension:

np.array([a[i,b[i]] for i in range(len(b))]

Upvotes: 0

Related Questions