CODEWITHSUNDEEP
pythonpython-3.xnumpy
OMRY VOLK
OMRY VOLK

Reputation: 1481

Access multiple columns in a 3D numpy array without looping

I have a large 3D array A with shape (N, M, L).

I have a list of coordinates of columns I want to access stored in a 2D array B:

[[i1 j1]
 [i2 j2]
 [i3 j3]
 ....   ]

I have something that works OK but involves looping over B and accessing A multiple times. Is there a way to avoid this using slicing or another method?

My code so far:

data_out = []
for p in B:
    i, j = p
    col = A[:, i, j]
    data_out.append(col)

Upvotes: 2

Views: 181

Answers (1)

Mad Physicist
Mad Physicist

Reputation: 114230

Use fancy indexing:

A[(slice(None), *B.T)].T

The explicit parentheses are necessary to use star expansion, which means that you have to write out : explicitly as slice(None). You can also do

A[:, B[:, 0], B[:, 1]].T

Upvotes: 2

Related Questions