按规则为DataFrame分配唯一ID：A类型逐行累加，其他类型同组同ID

给定如下DataFrame：

import pandas as pd

df = pd.DataFrame({'type': ['A', 'A', 'A', 'B', 'B', 'B', 'A', 'A','C','D','C','D','D','A', 'A'],
                  })

需要基于type列生成唯一ID，规则如下：

• 当type为*'A'*时，每行ID逐行累加
• 当type为其他值时，连续相同的type为同一组，组内所有行使用同一个ID

最终期望得到的结果：

df = pd.DataFrame({'type': ['A', 'A', 'A', 'B', 'B', 'B', 'A', 'A','C','D','C','D','D','A', 'A'],
                   'id': [1, 2, 3, 4, 4, 4, 5,6, 7, 8, 9, 10, 10, 11, 12],
                  })

可以通过以下步骤实现需求：

1. 标记出需要单独计数的'A'行，以及需要分组的非'A'连续组
2. 生成分组标识，再将分组标识转换为连续的ID

具体代码如下：

import pandas as pd

df = pd.DataFrame({'type': ['A', 'A', 'A', 'B', 'B', 'B', 'A', 'A','C','D','C','D','D','A', 'A'],
                  })

# 创建分组键：非'A'的连续相同type为一组，每个'A'单独为一组
group_key = (
    (df['type'] != df['type'].shift()) | (df['type'] == 'A')
).cumsum()

# 将分组键作为ID
df['id'] = group_key

print(df)

逻辑说明： #

• df['type'] != df['type'].shift()：识别非'A'类型的连续组起始行（当前行与上一行type不同时标记为True）
• | (df['type'] == 'A')：将所有'A'行都标记为新组起始行，确保每个'A'行单独成组
• .cumsum()：对布尔值累加生成连续的分组编号，直接作为ID即可

运行代码后输出结果与期望完全匹配。

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