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

例如,我有一个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]])]

有没有更好的方法来遍历它?谢谢。

回答:

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

以上是 如何根据numpy中的条件拆分数组? 的全部内容, 来源链接: utcz.com/qa/411993.html

回到顶部