使用pandas应用多个布尔掩码在列中设置值的有效方法

Question

我有几个表达式可以选择数据框中的某些行（df）并返回多个布尔数组，如果你愿意，可以使用掩码。我想在此数据框中添加一个新列，并根据这些掩码用简单的整数（“成员资格”指示符）填充它。在最简单的方法中，我将使用此代码：

df.assign(membership=3)
df[mask1].membership = 0
df[mask2].membership = 1
df[mask3].membership = 2

这应该产生一个列（df["membership"]），如下所示：

1    0
2    0
3    1
4    1
5    3
Name: membership, dtype: int64

掩码是由函数创建的（使用数据框中的其他列来创建输出）作为字典，如下所示：

maskDict = {0:mask1, 1:mask2, 2:mask3}

其中mask(1..3)是包含[True,False,...]等的大小相等的数组。

我知道可能没有必要改进代码。但是，在比较更多条件时，我将在字典中使用更多布尔数组。

将该功能应用于数据框并设置这些值的就地解决方案也将受到赞赏。

Answer 1

使用numpy.select：

df["membership"] = np.select([mask1, mask2, mask3], [0,1,2], default=3)

如果输入为dict：

df["membership1"] = np.select(list(maskDict.values()), list(maskDict.keys()), default=3)