python - Pandas 有条件地创建系列/数据框列

Question

如何将 color 列添加到以下数据框中，以便 color='green' if Set == 'Z',和 color='red' 否则？

    Type       Set
1    A          Z
2    B          Z           
3    B          X
4    C          Y

df = {'Type': ['A', 'B', 'B', 'C'], 'Set': ['Z', 'Z', 'X', 'Y']}

Answer 1

如果您只有两个选择:

df['color'] = np.where(df['Set']=='Z', 'green', 'red')

例如，

import pandas as pd
import numpy as np

df = pd.DataFrame({'Type':list('ABBC'), 'Set':list('ZZXY')})
df['color'] = np.where(df['Set']=='Z', 'green', 'red')
print(df)

产量

  Set Type  color
0   Z    A  green
1   Z    B  green
2   X    B    red
3   Y    C    red

---

如果您有两个以上的条件，请使用 np.select 。例如，如果您希望 color 为

• 黄色 当(df['Set'] == 'Z') & (df['Type'] == 'A')
• 否则 blue 当 (df['Set'] == 'Z') & (df['Type'] == 'B')
• 否则紫色当(df['Type'] == 'B')
• 否则黑色,

然后使用

df = pd.DataFrame({'Type':list('ABBC'), 'Set':list('ZZXY')})
conditions = [
    (df['Set'] == 'Z') & (df['Type'] == 'A'),
    (df['Set'] == 'Z') & (df['Type'] == 'B'),
    (df['Type'] == 'B')]
choices = ['yellow', 'blue', 'purple']
df['color'] = np.select(conditions, choices, default='black')
print(df)

产生

  Set Type   color
0   Z    A  yellow
1   Z    B    blue
2   X    B  purple
3   Y    C   black