Numpy的高级索引

Question

4

我有两个数组: 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)。我想知道是否还有其他解决方案？

- justanothercoder

1个回答

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- Divakar · Accepted Answer

您可以使用2D扩展范围数组作为与索引数组index形状对齐的第一轴，因此可以使用高级索引，如下所示 -

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