python - Numpy 中的高级索引

Question

我有两个数组:A 这样 A.shape = (N, M, K) 和 index 这样 index.shape = (N, M).因此，我想要得到这样的数组 B，即 B[i, j] = A[i, j, index[i, j]]。如何以最简单的方式实现？

我想到的一个解决方案是我们可以这样做:B = A.reshape(-1, K)[np.arange(N * M), index.reshape(-1)].reshape(N, M)。请问还有其他解决办法吗？

Answer 1

您可以使用 2D 扩展范围数组作为与索引数组 index 的形状对齐的第一个轴，因此使用 advanced-indexing , 像这样 -

B = A[np.arange(N)[:,None], np.arange(M), index]