在numpy数组中仅获取满足条件的值

Question

获取NumPy数组中的元素（甚至更有趣的是其中一部分）必须有一种（非常）快速和高效的方法。 假设我有一个NumPy数组：

import numpy as np
a = np.arange(-10,10)

现在，如果我有一个列表:

s = [9, 12, 13, 14]

我可以从a中选择元素：

a[s]  #array([-1,  2,  3,  4])

我该如何创建一个(numpy)数组，由a[s]中满足条件的元素组成，即为正数（或负数）？ 它应该得到以下结果：

np.ifcondition(a[s]>0, a[s])  #array([2,  3,  4])

这看起来很简单，但我找不到一个简洁明了的表达方式。我确定口罩可以，但对我来说似乎并不是非常直接的方法。 然而，也不是：

a[a[s]>0]
a[s[a[s]>0]]

实际上是很好的选择。

- gluuke

---

大多数像np.clip或np.where这样的工具会保留数组的原始大小，所以它们并不真正符合我的需求。 - gluuke

Answer 1

这样怎么样：

In [19]: b = a[s]

In [20]: b[b > 0]
Out[20]: array([2, 3, 4])

- unutbu

---

没错！...它似乎可以工作：我只需要定义一个新数组。 - gluuke

a[s][a[s] > 0]也可以运行，但会导致NumPy重新计算a[s]两次。 - unutbu

这是做这种事情的标准方式。（我给你点赞） - mgilson

我可以在条件中使用这个，但我也想设置一个包含b中所有元素的默认条件。我该如何设置？（我正在尝试将其作为参数传递给函数，并希望将默认设置为全部包括） - Parthiban Rajendran

Answer 2

你应该毫不犹豫地接受unutbu的答案，这是我在numpy中处理此类情况时通常使用的方法。但是考虑到有多种方法可以解决问题、需要在numpy之外工作的情况或者中间数组非常庞大的情况下，我会添加以下备选方案:

In [3]: [a[S] for S in s if a[S]>0]
Out[3]: [2, 3, 4]

再次强调，unutbu的方法速度显著更快。但我喜欢这种方法，因为它可以进一步简化。如果你有一个昂贵的奇怪函数（例如不是索引），并且想要测试该函数的结果，可以采用以下方法：

In [5]: [f for S in s for f in [FunkyFunction(a[S])] if f>0]
Out[5]: [2, 3, 4]

这里的奇怪之处在于你在另一个列表内部创建了一个列表，但是这个内部列表只包含一个项目。基本上，你正在将值保存到变量 f 中，然后使用该值两次：一次用于测试该值 (f>0)，一次将该值用于列表中如果测试通过。

- Mike

Answer 3

我知道这是一个老问题，不确定在提问时是否有下面的内容，但是因为我有类似的问题，所以这就是我找到适合我的方法：

np.extract(a[s] > 0, a[s])

- Fernando Margueirat

---

“extract”函数已经存在很长时间，至少在2009年就有了。由于“a[s]>0”是一个布尔数组，文档说明你的“extract”表达式等同于“a[s][a[s]>0]”，但这个表达式已经被拒绝了，因为它计算了两次“a[s]”，而你的表达式也是如此。 - Mike