基于其他列的条件，使用特定值添加列的Pandas操作

Question

我有以下的pandas数据框：

import pandas as pd
import numpy as np

d = {'age' : [21, 45, 45, 5],
     'salary' : [20, 40, 10, 100]}

df = pd.DataFrame(d)

我希望添加一个额外的列，名为“is_rich”，根据一个人的工资来确定他/她是否富裕。我找到了多种实现方法：

# method 1
df['is_rich_method1'] = np.where(df['salary']>=50, 'yes', 'no')

# method 2
df['is_rich_method2'] = ['yes' if x >= 50 else 'no' for x in df['salary']]

# method 3
df['is_rich_method3'] = 'no'
df.loc[df['salary'] > 50,'is_rich_method3'] = 'yes'

导致：

然而我不明白哪种方法是首选。根据您的应用程序，所有方法是否都同样好？

- Rutger Hofste

Answer 1

使用timeit， Luke！

结论
列表推导在处理较小数据量时表现最佳，因为它们几乎没有开销，即使它们不是向量化的。然而，在处理大量数据时，loc和numpy.where表现更好-矢量化获胜。

请记住，方法的适用性取决于您的数据、条件数量和列的数据类型。我的建议是在确定选项之前在您的数据上测试各种方法。

但是，在这里可以肯定的一点是，列表推导具有相当竞争力——它们是用C实现的，并且针对性能进行了高度优化。

---

参考基准测试代码。以下是正在计时的函数：

def numpy_where(df):
  return df.assign(is_rich=np.where(df['salary'] >= 50, 'yes', 'no'))

def list_comp(df):
  return df.assign(is_rich=['yes' if x >= 50 else 'no' for x in df['salary']])

def loc(df):
  df = df.assign(is_rich='no')
  df.loc[df['salary'] > 50, 'is_rich'] = 'yes'
  return df

- cs95

---

如果我希望在lis_comp的else子句中什么也不发生，我该怎么做？ - rish

Answer 2

另一种方法是使用 pandas 中的 mask 方法（取决于用例，也可以使用 where 方法）。首先初始化一个带有默认值（选择为 "no"）的 Series，并根据条件替换其中的一些值（有点像 loc[] 和 numpy.where() 的混合体）。

df['is_rich'] = pd.Series('no', index=df.index).mask(df['salary']>50, 'yes')

这可能是最快的选择。例如，对于一个有10万行的框架，mask()选项比loc选项快40%。¹ 我还更新了cs95答案中的perfplot基准测试，以比较mask方法与其他方法的性能表现：

---

¹: 这个基准测试结果比较了mask和loc。

def mask(df):
    return df.assign(is_rich=pd.Series('no', index=df.index).mask(df['salary']>50, 'yes'))

df = pd.DataFrame({'salary': np.random.rand(10_000_000)*100})

%timeit mask(df)
# 391 ms ± 3.87 ms per loop (mean ± std. dev. of 10 runs, 100 loops each)

%timeit loc(df)
# 558 ms ± 75.6 ms per loop (mean ± std. dev. of 10 runs, 100 loops each)

- cottontail