在不循環的情況下訪問 3D numpy 數組中的多列

Question

我有一個形狀為(N, M, L)的大型 3D 陣列A 。

我有一個列的坐標列表，我想訪問存儲在二維數組B ：

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

我有一些工作正常但涉及循環B並多次訪問A 。 有沒有辦法使用切片或其他方法來避免這種情況？

到目前為止我的代碼：

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

Answer 1

使用花哨的索引：

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

顯式括號是使用星型擴展所必需的，這意味着您必須明確寫出:為slice(None) 。 你也可以這樣做

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