CODEWITHSUNDEEP
pythonarraysnumpyslice
Faur
Faur

Reputation: 5968

How to slice a np.array using a list of unknown length?

I am probably using the wrong names/notation (and an answer probably exists on SO, but I can't find it). Please help me clarify so that I can update the post, and help people like me in the future.

I have an array A of unknown dimension n, and a list of indexes of unknown length l, where l<=n.

I want to be able to select the slice of A that correspond to the indexes in l. I.e I want:

A = np.zeros([3,4,5])
idx = [1,3]
B = # general command I am looking for

B_bad = A[idx] # shape = (2,4,5), not what I want!
B_manual = A[idx[0], idx[1]] # shape = (5), what I want, but as a general expression.
# it is okay that the indexes are in order i.e. 0, 1, 2, ...

Upvotes: 1

Views: 308

Answers (1)

nicoco
nicoco

Reputation: 1553

You need a tuple:

>>> A[tuple(idx)]
array([0., 0., 0., 0., 0.])
>>> A[tuple(idx)].shape
(5,)

Indexing with a list doesn't have the same meaning. See numpy's documentation on indexing for more information.

Upvotes: 2

Related Questions