Pandas技巧：基于条件筛选创建新列并分配递增整数值

针对你需求的场景，我来分享个简洁优雅的实现方式，完美解决你提到的两个问题：

import pandas as pd

# 初始化你的DataFrame
df = pd.DataFrame({
    'id': [1,2,3,4,5],
    'category': ['A','A','B','C','A']
})

# 一行代码生成目标value列
df['value'] = df.groupby(df['category'].eq('A'))['category'].cumcount().add(1).where(df['category'].eq('A'), pd.NA).astype('Int64')

代码细节解释： #

• df['category'].eq('A')：生成布尔掩码，标记出所有category为'A'的行
• groupby(...)['category'].cumcount()：基于布尔掩码分组，对True组（也就是'A'类行）生成从0开始的连续计数
• .add(1)：把计数从0起始改成1起始，匹配你需要的1、2、3...递增序列
• .where(df['category'].eq('A'), pd.NA)：只保留'A'类行的计数，其他行替换为Pandas的标准空值标记pd.NA
• .astype('Int64')：转换为支持空值的整数类型（注意是大写I的Int64），既保留了整数类型，又能兼容空值，彻底解决原来的浮点型问题

最终输出效果： #

id category  value
0   1        A      1
1   2        A      2
2   3        B   <NA>
3   4        C   <NA>
4   5        A      3

这个实现的优势：

• 无需预先初始化NaN列，代码更简洁紧凑
• 直接生成支持空值的整数列，彻底规避浮点型问题
• 完全利用Pandas原生方法，符合Pythonic的简洁风格

内容的提问来源于stack exchange，提问作者skinnyas123