CODEWITHSUNDEEP
pythonnumpy
user357269
user357269

Reputation: 1913

slicing by indices on multiple axes numpy

A = np.arange(120).reshape(2, 3, 4, 5)
is_ = [1, 2]
js = [0, 1, 2, 3]

A[:, :, is_, :][:, :, :, js].shape == (2, 3, 2, 4)

Is there a better way of doing the double slice here?

I tried A[:, :, is_, js] but that does it "zip" style.

Efficiency would be nice too, I'm having to do this in a double loop...

Upvotes: 1

Views: 538

Answers (1)

yatu
yatu

Reputation: 88226

You can do it in a single indexing step. You just need to add a new axis to either of the indexing arrays so they are broadcastable:

is_ = np.array([1, 2])
js = np.array([0, 1, 2, 3])

A[:, :, is_[:,None], js]

Upvotes: 2

Related Questions