如何从一个新的数组中提取所有零和非零值？

Question

给出一个numpy数组：

arr = [0,0,0,0,0,2,0]

如何在新数组中提取所有零元素：

new_arr = [0,0,0,0,0,0]

非零元素：

new_arr_2 = [2]

我尝试通过将数组转换为列表并使用for循环进行尝试。另外，我尝试了：

np.argwhere(arr == 0).ravel()

但是，我只是获取索引。关于如何使用numpy提取零，非零甚至特定数字值的任何想法？

Answer 1

通过使用numpy.where，可以获得满足条件的项的索引并获得相应的值：

import numpy as np

arr = np.array([0,0,0,0,0,2,0])

new_arr_0 = arr[np.where(arr==0)]
new_arr_no_0 = arr[np.where(arr!=0)]

当然，按照@Sheldore和@DeveshKumarSingh的建议，您可以简单地选择：

new_arr_0 = arr[arr==0]
new_arr_no_0 = arr[arr!=0]

对于小数组，您不应在计算速度上发现大的差异。

Answer 2

您可以使用列表推导代替numpy：

arr = [0,0,0,0,0,2,0]
new_arr = [a for a in arr if a == 0]
new_arr_no_zero = [a for a in arr if a != 0]