How to extract specific rows from an nd array in python

Question

I have an array b which is of the shape E x B x 3. I have another array a which specifies which 3 elements to take our of b.

The following code works (in this example E=2, B=4):

import numpy as np

a = [1, 1, 0, 0]
b = np.array([[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 3.0]],
            [[2.0, 0.0, 0.0], [0.0, 2.0, 0.0],  [0.0, 0.0, 2.0], [0.0, 0.0, 3.0]]])
# n_pred = np.transpose(n_pred, axes=[1, 0, 2])
c = []
for i, idx in enumerate(a):
    c.append(b[idx, i])
c = np.array(c)
print(c)

My question is, is there a more efficient way to do this? (maybe using some built-in numpy function?

jpp · Accepted Answer · 2018-07-23 17:42:52Z

1

You can index by the first two dimensions:

c = b[a, range(len(a))]

print(c)

array([[ 2.,  0.,  0.],
       [ 0.,  2.,  0.],
       [ 0.,  0.,  1.],
       [ 0.,  0.,  3.]])

edited Jul 23, 2018 at 17:42

answered Jul 23, 2018 at 12:57

jpp

166k37 gold badges301 silver badges362 bronze badges

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

1 Comment

itzik Ben Shabat Over a year ago

Thanks! I knew it was something super simple that I was missing.

Collectives™ on Stack Overflow

How to extract specific rows from an nd array in python

1 Answer 1

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related