如何根据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

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