Pandas中按指定列分组生成1到n递增序号新列的高效方法

需求说明 #

基于如下示例DataFrame，需新增列c：按列a的取值分组，组内序号从1开始逐行递增到当前组的总行数，替代性能较差的for循环实现，适配大数据量处理场景。

示例原始数据构造代码：

import pandas as pd  

df = pd.DataFrame({'a': [1,1,1,2,2,2,2,2,3,3,3,3,4,4,4,4,4,4],
                   'b': [3,4,3,7,5,9,4,2,5,6,7,8,4,2,4,5,8,0]})

预期输出效果：

a   b   c
0   1   3   1
1   1   4   2
2   1   3   3
3   2   7   1
4   2   5   2
5   2   9   3
6   2   4   4
7   2   2   5
8   3   5   1
9   3   6   2
10  3   7   3
11  3   8   4
12  4   4   1
13  4   2   2
14  4   4   3
15  4   5   4
16  4   8   5
17  4   0   6

实现代码 #

直接使用Pandas内置的分组累计计数方法cumcount()即可，该方法为底层向量化实现，性能极高：

# cumcount默认从0开始计数，统一加1即可匹配从1开始的需求
df['c'] = df.groupby('a').cumcount() + 1

补充说明 #

• 该实现无Python层面的循环开销，处理速度比手写for循环快2~3个数量级，完全适配亿级以内的大数据量处理场景
• 方法兼容性强，Pandas 0.20及以上正式版本均可直接调用，无需安装额外依赖
• 若需要组内按其他字段排序后再生成序号，只需在分组前对数据做排序即可，例：组内按b列升序生成序号的代码如下

df['c'] = df.sort_values(by=['a', 'b']).groupby('a').cumcount() + 1

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