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

2024-03-29 11:35:56 发布

您现在位置:Python中文网/ 问答频道 /正文

例如,我有一个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并创建两个新数组。但我想知道纽比能提供更好的方法吗?

类似地,对于多维数组,例如

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

有没有更好的方法来代替它?谢谢。


Tags: 方法ltgtnp数字数组arrayndarray
2条回答
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]])]

这可能是一个快速的解决方案

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

相关问题 更多 >