Python 如何设置nan值为numpy数组多条件

发布于03月16日

我有一些数值的NumPy数组，我希望数值为2、3和7的点将保持不变，但其他值将是NaN.

原始的nyopartist数组看起来像这样

  [[0, 7, 2, ..., 1, 6, 0],
   [5, 3, 0, ..., 2, 5, 0],
   ...,
   [8, 2, 0, ..., 1, 0, 7],
   [0, 3, 2, ..., 0, 2, 0],
   [7, 1, 0, ..., 0, 7, 0]],

我的预期结果是

  [[nan, 7, 2, ..., nan, nan, nan],
   [nan, 3, nan, ..., 2, nan, nan],
   ...,
   [nan, 2, nan, ..., nan, nan, 7],
   [nan, 3, 2, ...,   nan, 2, nan],
   [7, nan, nan, ..., nan, 7, nan]],

推荐答案

结合np.isin和np.where:

a = np.array([[0, 7, 2, 1, 6, 0],
              [5, 3, 0, 2, 5, 0],
              [8, 2, 0, 1, 0, 7],
              [0, 3, 2, 0, 2, 0],
              [7, 1, 0, 0, 7, 0]],)

out = np.where(np.isin(a, [2, 3, 7]), a, np.nan)

如果a有一个float dtype，你也可以修改:

a = np.array([[0, 7, 2, 1, 6, 0],
              [5, 3, 0, 2, 5, 0],
              [8, 2, 0, 1, 0, 7],
              [0, 3, 2, 0, 2, 0],
              [7, 1, 0, 0, 7, 0]], dtype=float)

a[~np.isin(a, [2, 3, 7])] = np.nan

输出:

array([[nan,  7.,  2., nan, nan, nan],
       [nan,  3., nan,  2., nan, nan],
       [nan,  2., nan, nan, nan,  7.],
       [nan,  3.,  2., nan,  2., nan],
       [ 7., nan, nan, nan,  7., nan]])