基于numpy python中的条件子集二维数组

时间:2016-05-12 05:04:41

标签: python numpy

我有一个大小为3600 * 7200的numpy 2D数组。我有另一个相同形状的数组,我想用作掩码。

问题在于当我做这样的事情时:

import numpy as np
N = 10
arr_a = np.random.random((N,N))
arr_b = np.random.random((N,N))
arr_a[arr_b > 0.0]

生成的数组不再是2D,而是1D。如何获得2D数组呢?

1 个答案:

答案 0 :(得分:2)

您可以使用np.where来保留形状:

np.where(arr_b > 0.0, arr_a, np.nan)

当arr_b的值大于0时,它将从arr_a获取相应的值,否则它将使用np.nan。

import numpy as np
N = 5
arr_a = np.random.randn(N,N)
arr_b = np.random.randn(N,N)
np.where(arr_b > 0.0, arr_a, np.nan)

Out[107]: 
array([[ 0.5743081 ,         nan, -1.69559034,         nan,  0.4987268 ],
       [ 0.33038264,         nan, -0.27151598,         nan, -0.73145628],
       [        nan,  0.46741932,  0.61225086,         nan,  1.08327459],
       [        nan, -1.20244926,  1.5834266 , -0.04675223, -1.14904974],
       [        nan,  1.20307104, -0.86777899,         nan,         nan]])