如何根据numpy中的条件拆分数组？

Question

例如，我有一个 ndarray 即：

 a = np.array([1, 3, 5, 7, 2, 4, 6, 8])

现在我想将 a 分成两部分，一个是所有数字 <5，另一个是所有 >=5：

 [array([1,3,2,4]), array([5,7,6,8])]

当然我可以遍历 a 并创建两个新数组。但我想知道 numpy 是否提供了一些更好的方法？

同样，对于多维数组，例如

array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9],
       [2, 4, 7]])

我想根据第一列 <3 和 >=3 拆分它，结果是：

 [array([[1, 2, 3],
        [2, 4, 7]]),
 array([[4, 5, 6],
        [7, 8, 9]])]

有没有更好的方法而不是遍历它？谢谢。

原文由 Clippit 发布，翻译遵循 CC BY-SA 4.0 许可协议

Answer 1

import numpy as np

def split(arr, cond):
  return [arr[cond], arr[~cond]]

a = np.array([1,3,5,7,2,4,6,8])
print split(a, a<5)

a = np.array([[1,2,3],[4,5,6],[7,8,9],[2,4,7]])
print split(a, a[:,0]<3)

这会产生以下输出：

 [array([1, 3, 2, 4]), array([5, 7, 6, 8])]

[array([[1, 2, 3],
       [2, 4, 7]]), array([[4, 5, 6],
       [7, 8, 9]])]

原文由 NPE 发布，翻译遵循 CC BY-SA 3.0 许可协议

Answer 2

一维数组 a = numpy.array([2,3,4,…]) a_new = a[(a < 4)] # 获取小于 5 的元素

基于列的二维数组（考虑第 i 列的值应小于 5， a = numpy.array([[1,2],[5,6],…] a = a[(a[:,i ] < 5)]

如果您的条件是基于多列的，那么您可以创建一个新数组，将条件应用于列。然后您可以将新数组与值 5（根据我的假设）进行比较以获取索引并遵循上述代码。请注意，无论我在 () 中写入什么，都会返回索引数组。

原文由 Apurba A 发布，翻译遵循 CC BY-SA 4.0 许可协议