如何在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]]

在这种情况下，还有更好的方法吗？

- f. c.

---

你为什么要使用argwhere？为什么不直接使用xx[xx>3]？ - user2357112

这只是一个例子。在我的情况下，我确实有类似于示例中的索引数组index。 - f. c.

xx[tuple(index)] 对我来说可行。 - Karl Knechtel

@KarlKnechtel：看起来它是工作的，但只是偶然，因为“index”碰巧是它自己的转置。如果您尝试使用“index = np.argwhere（xx> 2）”，则会出现IndexError。 - user2357112

@f.c. 这个答案对你有用吗？ - j-i-l

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]]

当使用nonzero时，它会返回一个元组：

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

请注意，后者是索引所需的内容。此外，后者几乎是前者的转置，因此，如果您被一个由argwhere产生的数组格式所困扰，您可以转置并将其转换为元组，并且就可以了。所以在您的情况下是 tuple(np.transpose(index))。 如果您正在使用argwhere来获取索引，则建议改用nonzero。

- j-i-l

---

argwhere 只是返回 np.transpose(np.nonzero(...))，将包含 n 个数组的元组转换为 n 列数组。 - hpaulj

@hpaulj 是的，完全正确！这就是我想用 argwhere 按元素分组，而 nonzero 按维度分组的意思。这也是为什么你可以直接使用 nonzero 的输出进行索引，但需要先转置 argwhere 的输出（或者直接使用 nonzero）。 - j-i-l

np.transpose(index) 不足够，你仍然需要转换回一个元组(tuple(index.T))，因为一个二维数组并不等同于由一维数组组成的元组用作索引。 - user2357112

@user2357112supportsMonica 啊哦，谢谢你指出这个问题，我忘记检查返回类型了。 - j-i-l