如何快速从numpy数组中获取特定索引？

Question

但是我没有索引值，我只有另一个数组中相同索引位置的数字。例如，我有

a = array([3,4,5,6])
b = array([0,1,0,1])

有没有一种NumPy方法可以快速查看它们并从a中提取所有索引与b中所有1的索引匹配的值？ 我期望结果为：

array([4,6])

值得一提的是，我的a数组是多维的，而我的b数组始终只有0或1的值。我尝试使用NumPy的logical_and函数，但是当a和b具有不同的维度时，它会返回ValueError：

a = numpy.array([[3,2], [4,5], [6,1]])
b = numpy.array([0, 1, 0])
print numpy.logical_and(a,b)

ValueError: operands could not be broadcast together with shapes (3,2) (3,) 

虽然这种方法似乎在a是平面的情况下有效，但无论如何，numpy.logical_and（）的返回类型都是布尔值，而我不想要这个。还有其他的方法吗？再次强调，在上面的第二个示例中，期望的返回值将是

array([[4,5]])

显然我可以写一个简单的循环来完成这个任务，但我只是想要更简洁的方法。

编辑：

这会引入更多的限制，我还应该提到多维数组a的每个元素可能具有任意长度，而不匹配其相邻元素。

- pretzlstyle

---

a[b.astype(bool)] 怎么样？ - Divakar

我刚试了一下，谢谢。这很棒。不幸的是，它对于多维数组a的情况不起作用 :( - pretzlstyle

添加一个多维数组的示例案例？ - Divakar

Answer 1

您可以简单地使用高级索引。

b == 1

将会给你一个布尔型数组：

>>> from numpy import array
>>> a = array([3,4,5,6])
>>> b = array([0,1,0,1])
>>> b==1
array([False,  True, False,  True], dtype=bool)

您可以将其作为索引传递给a。

>>> a[b==1]
array([4, 6])

您的第二个示例的演示:

>>> a = array([[3,2], [4,5], [6,1]])
>>> b = array([0, 1, 0])
>>> a[b==1]
array([[4, 5]])

- timgeb

---

我非常喜欢这个，但它似乎不能处理任意多维情况。 - pretzlstyle

@jphollowed 对于您的第二个例子[[3,2], [4,5], [6,1]]和[0, 1, 0]它是有效的，您有哪些困难的情况？ - timgeb

@ah 你说得对，我犯了一个错误。这是最优雅的答案。 - pretzlstyle

如果你发现这是最容易处理的方法，那么确实这是一个好方法，并且应该被接受。值得注意的是，它需要首先创建一个新的临时布尔数组，但除非你正在使用大规模数组，否则这不应该是一个问题。还要记得将布尔数组放在正确的轴上！例如使用a[:, b==1]从轴1中获取，而不是轴0中获取。 - Alex Riley

我应该补充说明，从我的实验来看，如果你的对象之一是Python列表，而另一个是NumPy数组，那么这将无法工作。 - pretzlstyle

@jphollowed，您可以将该列表转换为array。它不能与Python列表lst一起使用的原因是lst == 1只会返回False，而不是布尔数组。 - timgeb

Answer 2

您可以使用 compress 来实现：

>>> a = np.array([3,4,5,6])
>>> b = np.array([0,1,0,1])
>>> a.compress(b)
array([4, 6])

您可以为多维情况提供一个 axis 参数：

>>> a2 = np.array([[3,2], [4,5], [6,1]])
>>> b2 = np.array([0, 1, 0])
>>> a2.compress(b2, axis=0)
array([[4, 5]])

即使你要索引的a轴与b不同，这种方法仍然有效。

- Alex Riley