如何根据条件在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]])]

有没有比遍历更好的方法呢？谢谢。

Answer 1

一维数组

a = numpy.array([2,3,4,...]) 这里我们创建了一个一维数组，里面有一些数字。接下来，a_new = a[(a < 4)] 这行代码的意思是，我们想要从这个数组中找出所有小于4的数字，结果会存放在新的数组a_new里。

二维数组，按列处理（假设我们关注的列的值应该小于5）

a = numpy.array([[1,2],[5,6],...]) 这里我们创建了一个二维数组，里面有几行几列的数字。然后，a = a[(a[:,i] < 5)] 这行代码的意思是，我们要检查第i列的所有数字，找出那些小于5的数字，并把符合条件的行保留下来。

如果你的条件涉及多列，那么你可以创建一个新的数组，应用这些条件在不同的列上。之后，你只需要把这个新数组和5进行比较（根据我的假设），就能找到符合条件的索引，然后可以使用上面的代码。请注意，括号里的内容会返回一个索引数组。

Answer 2

这可能是个快速的解决办法。

a = np.array([1,3,5,7])
b = a >= 3 # variable with condition
a[b] # to slice the array
len(a[b]) # count the elements in sliced array

Answer 3

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