代码看起来"not pythonic" - 嵌套的 np.where() 将列添加到 pd.dataframe

Question

我有一个pd数据帧，其中包括CompTotal和CompFreq列。

我想添加第三列-NormalizedAnnualCompensation，它使用以下逻辑如果CompFreq为Yearly，则使用CompTotal中的现有值如果CompFreq为Monthly，则将CompTotal中的值乘以12如果CompFreq为Weekly，则将CompTotal中的值乘以52

我最终使用np.where()基本上编写了一个嵌套的if语句，就像我习惯于在excel中组合在一起一样(对一般的编码来说很陌生(——如下所示。

我的问题是——我能做得更好吗？根据我所读到的内容和迄今为止所学到的知识，这感觉不太像蟒蛇。

df['NormalizedAnnualCompensation'] = np.where(df['CompFreq']=='Yearly',df.CompTotal, 
(np.where(df['CompFreq']=='Monthly', df.CompTotal * 12,
(np.where(df['CompFreq']=='Weekly',df.CompTotal *52,'NA')
))))

提前谢谢。

Answer 1

；适当的"；做事的方法，所以你已经得到了正确的！

尽管如此，你还是可以通过询问不同的方法来学习(尽管这可能超出了stackoverflow的意图(。

例如，您可以考虑仅通过使用掩码并仅访问要设置的数据帧的某个特定区域(pd.DataFrame.loc(来使用panda：

df["NormalizedAnnualCompensation"] = "NA"
mask = df["CompFreq"]=="Yearly"
df.loc[mask, "NormalizedAnnualCompensation"] = df.loc[mask, "CompTotal"]
mask = df["CompFreq"]=="Monthly"
df.loc[mask, "NormalizedAnnualCompensation"] = df.loc[mask, "CompTotal"] * 12
mask = df["CompFreq"]=="Weekly"
df.loc[mask, "NormalizedAnnualCompensation"] = df.loc[mask, "CompTotal"] * 52

如果你真的只想比较该列的相等性，并且每种情况都填充一个固定值(即CompTotal是整个数据帧的常数，你可以考虑简单地使用pd.Series.map，比较以下实现类似效果的最小示例：

In [1]: pd.Series(np.random.randint(4, size=10)).map({0: "zero", 1: "one", 2: "two"}).fillna(
...:     "NA"
...: )
Out[1]:
0      NA
1     two
2      NA
3    zero
4     two
5    zero
6     one
7     two
8      NA
9     two
dtype: object

Answer 2

np.where()适用于简单的if-then-else处理。然而，如果您有多个条件要测试，嵌套np.where()看起来会很复杂，很难阅读。在这种情况下，您可以使用np.select()获得更干净、更可读的代码，如下所示：

condlist   = [df['CompFreq']=='Yearly', df['CompFreq']=='Monthly', df['CompFreq']=='Weekly']
choicelist = [df.CompTotal,             df.CompTotal * 12,         df.CompTotal * 52]
df['NormalizedAnnualCompensation'] = np.select(condlist, choicelist, default='NA')