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

Question

我有一个A带有 shape的大型 3D 数组(N, M, L)。

我有一个要访问的列的坐标列表，存储在二维数组中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)

标签： pythonpython-3.xnumpy

Answer 1

使用花哨的索引：

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

显式括号是使用星形扩展所必需的，这意味着您必须:显式写为slice(None). 你也可以做

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