在不循环的情况下访问 3D numpy 数组中的多列

Question

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)

Answer 1

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