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, Apr 15 '20 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, Apr 15 '20 at 12:57
Does this answer your question? [NumPy selecting specific column index per row by using a list of indexes](https://stackoverflow.com/questions/23435782/numpy-selecting-specific-column-index-per-row-by-using-a-list-of-indexes) — Nuageux, Apr 15 '20 at 13:25

score 2 · Answer 1 · answered Apr 15 '20 at 13:27

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 '20 at 13:27

Nils Werner

34,832
7
76
98

Take a look at the newish `np.take_along_axis` – hpaulj Apr 15 '20 at 16:53

sltzgs · Answer 2 · 2020-04-15T13:25:43.583

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 '20 at 13:25

answered Apr 15 '20 at 13:14

sltzgs

166
6

Thanks, this is indeed the desired action! Is there a way to do it without loops? – user44136 Apr 15 '20 at 13:15

hpaulj · Answer 3 · 2020-04-15T16:47:02.127

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.

Numpy matrix: select column based on list

4 Answers4