pandas 用 case_when 根据年龄进行条件替换

看过来

《pandas 教程》持续更新中，提供建议、纠错、催更等加作者微信: gairuo123（备注：pandas教程）和关注公众号「盖若」ID: gairuo。跟作者学习，请进入 Python学习课程。欢迎关注作者出版的书籍：《深入浅出Pandas》和《Python之光》。

pandas 2.2 新增的 case_when() 条件替换让我们针对多个条件进行替换数据更加方便了，我们第一时间进行了教程的更新，并推出本案例看看怎么让它用起来。

注：本例在 pandas 2.2 下通过，同时建议安装 pyarrow 库。

数据与需求

我们构造源数据如下：

import pandas as pd

data = {'name': ['Jason', 'Molly', 'Tina', 'Jake', 'Amy'], 
        'age': [10, 52, 36, 24, 73]}
df = pd.DataFrame(data, columns = ['name', 'age'])
df
'''
    name  age
0  Jason   10
1  David   52
2  James   36
3   Jake   24
4   Mark   73
'''

我们需要增加一列 elderly，根据年龄列对每个人进行标注，逻辑用 SQL 伪代码表示大致如下：

if age < 10 then baby
if age >= 10 and age < 20 then kid 
if age >=20 and age < 30 then young 
if age >= 30 and age < 50 then mature 
if age >= 50 then grandpa

代码

本例非常适合用 case_when() 条件替换，代码如下：

elderly = df.age.case_when([
    (df.age.lt(10),                 'baby'),
    (df.age.ge(10) & df.age.lt(20), 'kid'),
    (df.age.ge(20) & df.age.lt(30), 'young'),
    (df.age.ge(30) & df.age.lt(50), 'mature'),
    (df.age.ge(50),                 'grandpa'),
])
elderly
'''
0        kid
1    grandpa
2     mature
3      young
4    grandpa
Name: age, dtype: object
'''

df.assign(elderly=elderly)
'''
    name  age  elderly
0  Jason   10      kid
1  David   52  grandpa
2  James   36   mature
3   Jake   24    young
4   Mark   73  grandpa
'''

这样就得到了最终的结果。

（完）

pandas 用 case_when 根据年龄进行条件替换

数据与需求

代码

相关内容