如何用numpy中的索引索引数组？

Question

例如，我有一个数组及其索引：

  xx = np.arange(6).reshape(2,3)
  index = np.argwhere(xx>3)

如果我使用

 xx[index]

我收到一个错误：

IndexError: index 2 is out of bounds for axis 0 with size 2

但如果我写：

xx[index[0], index[1]]

我得到：

array([4, 5])

这正是我想要的。但是如果我有一个 N 维数组，我是否需要写

xx[index[0], index[1], ..., index[N]]

在这种情况下，还是有更好的方法来做到这一点？

Answer 1

从numpy 的文档中：

argwhere的输出不适合索引数组。为此，请改用 nonzero(a)。

所以，这将是推荐的方式，它工作得很好：

>>> xx = np.arange(6).reshape(2,3)
>>> index = np.nonzero(xx>3)
>>> xx[index]
array([4, 5])

的点np.argwhere实质上是它将组索引由元件和因此产生的输出将是形状的(M, xx.ndim)，其中M是非零元素的数量，这使得它显然不适合索引。

np.nonzero，另一方面，按维度对索引进行分组，为每个维度返回一个您需要的元组！

因此，基本上如果[x1,y1], [x2,y2]和[x3, y3]处的元素满足您的条件，argwhere将返回如下数组：

[[x1, y1], [x2, y2], [x3, y3]]

whilenonzero将返回一个元组：

([x1, x2, x3], [y1, y2, y3])

请注意，后者是您编制索引所需要的。另外，后者也几乎是前者的转置，所以如果你因为某种原因被argwhere你产生的形式数组卡住了，你可以转置并将其转换为元组，你就是金子。所以tuple(np.transpose(index))在你的情况下。如果您argwhere用于获取索引，那么我建议您直接使用nonzero。