Pandas 根據其他列的條件添加列

Question

我有以下 pandas dataframe：

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'

導致：

但是我不明白首選的方式是什么。 根據您的應用，所有方法是否都一樣好？

Answer 1

使用timeits ，盧克！

結論
列表推導在較小數據量上表現最佳，因為即使它們沒有矢量化，它們也會產生很少的開銷。 OTOH，對於更大的數據， 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

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 mil 行的幀， 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)