基於numpy python中的條件對二維數組進行子集化

Question

我有一個大小為 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]

結果數組不再是二維的，而是一維的。 我如何得到一個二維數組作為回報？

Answer 1

您可以使用 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]])