python - 在 np.where() 逻辑中指定 NaN 编码

Question

我的数据如下所示:

    id    case2_q6
0   300   3.0
1   304   4.0
2   306   3.0
3   309   1.0
4   311   3.0
5   312   4.0
6   314   NaN
7   315   2.0
8   316   3.0
9   317   3.0

并使用此 np.where() 函数调用来生成新变量:

df['fluid_2'] = np.where((df['case2_q6'] == 1) | (df['case2_q6'] == 2), 1, 0)

现在df具有fluid_2列，如下所示:

    id    case2_q6  fluid_2
0   300   3.0       0
1   304   4.0       0
2   306   3.0       0
3   309   1.0       1
4   311   3.0       0
5   312   4.0       0
6   314   NaN       0
7   315   2.0       1
8   316   3.0       0
9   317   3.0       0

如您所见，索引 6 处的 NaN 值已转换为 0。有没有办法设置 np.where() 以便在 fluid_2 中将它们保留为 NaN 值？

期望的输出是:

    id    case2_q6  fluid_2
0   300   3.0       0
1   304   4.0       0
2   306   3.0       0
3   309   1.0       1
4   311   3.0       0
5   312   4.0       0
6   314   NaN       NaN
7   315   2.0       1
8   316   3.0       0
9   317   3.0       0

保留 NaN 的位置。

Answer 1

可能的解决方案:

df['fluid_2'] = np.where(
    df['case2_q6'].isna(), np.nan, 
    np.where((df['case2_q6'] == 1) | (df['case2_q6'] == 2), 1, 0))

---

另一种可能的解决方案:

df['fluid_2'] = df['case2_q6'].clip(upper=1).mul(df['case2_q6'].isin([1,2]))

---

输出:

    id  case2_q6  fluid_2
0  300       3.0      0.0
1  304       4.0      0.0
2  306       3.0      0.0
3  309       1.0      1.0
4  311       3.0      0.0
5  312       4.0      0.0
6  314       NaN      NaN
7  315       2.0      1.0
8  316       3.0      0.0
9  317       3.0      0.0