Pandas基于ID分组和A列值新增标记列B的实现方法

问题场景 #

初始DataFrame结构如下：

需要新增列B完成分组标记，规则为：同一ID分组下，若A列的两个值相等且不为0则标记为2，否则标记为1，目标结果结构：

原有代码问题 #

你写的代码逻辑存在本质缺陷：多次通过过滤+concat拼接的方式只保留了满足B=2条件的行，直接丢弃了其余需要标记为1的行，自然无法返回完整的结果表。

正确实现代码 #

核心思路是用groupby + transform把分组判断结果广播到组内每一行，全程不丢弃原表数据，直接完成列赋值：

import pandas as pd

# 按ID分组生成判断条件，结果自动对齐原表每一行
# 组内满足两个条件则返回True：1.A列所有值相等 2.值不为0
match_mask = df.groupby('ID')['A'].transform(
    lambda col: (col.nunique() == 1) and (col.iloc[0] != 0)
)
# 按条件映射为目标标记值
df['B'] = match_mask.map({True: 2, False: 1})

逻辑说明 #

• transform方法会把分组计算的结果按照原表索引映射回每一行，不会修改原表的行数和顺序
• col.nunique() == 1用来判断当前ID组内A列是否所有值都相等
• 因为已经判定组内值全部相等，只需要取第一个值判断是否非0即可，效率高于遍历整列判断
• 最终通过布尔值映射直接得到1、2标记，逻辑简洁没有冗余的表拼接操作

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