实用的 Pandas 根据条件新建列的方法

工作中经常会碰到根据条件创建新的列,新建列的方法有很多,不同方法之间的效率差别很大,今天我们来看看比较常用的高效的根据条件来用 pandas 新建列。

通常我们遇到需要新建列的情况无外乎以下几种:

  1. 根据某一列新建另一列。
  2. 根据多列新建另一列。

接下来我们来看看以上两种情况该怎处理。

1.如何根据某一列新建另一列

In [2]: df = pd.DataFrame({'Distance': [1, 5, 40, 15]})
   ...: df
Out[2]:
   Distance
0         1
1         5
2        40
3        15

现在我们要新建一列对 Distance 列的数据进行分类,大于 10 的分类为 Long,小余2的分类为 Short,其它的为 Average。有多种方法可以实现该目的,可以用 lambda 函数、.loc、np.where、np.select 等。

# 用 lambda 实现
df['Category'] = df['Distance'].map(lambda x: 'Long' if x > 10 else ('Short' if x < 2 else 'Average'))

# 用 .loc 实现
df.loc[df['Distance']>10, 'Category'] = 'Long'
df.loc[df['Distance']<2, 'Category'] = 'Short'
df.loc[(df['Distance']>=2)&(df['Distance']<=10), 'Category'] = 'Average'

# 用 np.where 实现
df['Category'] = np.where(df['Distance']>10, 'Long', np.where(df['Distance']<2, 'Short', 'Average'))

# 用 np.select 实现
df['Category'] = np.select([df.Distance < 2, df.Distance > 10], ['short', 'long'], 'average')

2.如何根据多列新建另一列

In [3]: df = pd.DataFrame({'Type':list('ABBC'), 'Set':list('ZZXY')})

In [4]: df
Out[4]:
  Type Set
0    A   Z
1    B   Z
2    B   X
3    C   Y

现在要新建一列 color ,规则如下:

  • if Set 为 Z 且 Type 为 A,则 color 为 yellow
  • elif Set 为 Z 且 Type 为 B,则 color 为 blue
  • elif Type 为 B,则 color 为 purple
  • else color 为 black

我们可以用 np.select 和 自定义函数来解决上述问题。

# 用 np.select 实现
In [5]: 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)
  Type Set   color
0    A   Z  yellow
1    B   Z    blue
2    B   X  purple
3    C   Y   black

# 用自定义函数实现
def set_color(row):
    if (row['Set'] == 'Z') & (row['Type'] == 'A'):
        return 'yellow'
    elif (row['Set'] == 'Z') & (row['Type'] == 'B'):
        return 'blue'
    elif row['Type'] == 'B':
        return 'purple'
    else:
        return 'black'
    
df['color'] = df.apply(set_color, axis=1)

赶紧去试试吧!

为了及时收到文章更新,欢迎关注个人公众号:Python实战数据分析

发表评论

电子邮件地址不会被公开。