Numpy matrix: select column based on list

Question

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)

For each row in a you select the element at the column defined by b — user44136
– user44136, Commented Apr 15, 2020 at 12:56
You thus want the 1st element of the first row of a, the 0th element of the second row and the 2nd element of the last row (as defined in b) — user44136
– user44136, Commented Apr 15, 2020 at 12:57
Does this answer your question? NumPy selecting specific column index per row by using a list of indexes — Nuageux
– Nuageux, Commented Apr 15, 2020 at 13:25

Nils Werner · Accepted Answer · 2020-04-15 13:27:37Z

2

If you remove the extraneous dimensions from b

b = np.sqeeze(b)

You can use the following:

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

answered Apr 15, 2020 at 13:27

Nils Werner

37.2k7 gold badges85 silver badges108 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

hpaulj Over a year ago

Take a look at the newish np.take_along_axis

sltzgs · Accepted Answer · 2020-04-15 13:25:43Z

1

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))]]

edited Apr 15, 2020 at 13:25

answered Apr 15, 2020 at 13:14

sltzgs

1661 silver badge6 bronze badges

1 Comment

user44136 Over a year ago

Thanks, this is indeed the desired action! Is there a way to do it without loops?

hpaulj · Accepted Answer · 2020-04-15 16:47:02Z

1

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.

1 Comment

user44136 Over a year ago

This is indeed what I needed. Thanks!

Serge Ballesta · Accepted Answer · 2020-04-15 13:25:21Z

0

You could use a list comprehension:

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

answered Apr 15, 2020 at 13:25

Serge Ballesta

150k13 gold badges137 silver badges267 bronze badges

Collectives™ on Stack Overflow

Numpy matrix: select column based on list

4 Answers 4

1 Comment

1 Comment

1 Comment

Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

4 Answers 4

1 Comment

1 Comment

1 Comment

Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related