R中seq_len的Python等价实现：DataFrame分组内添加索引

问题描述 #

给定如下DataFrame：

Col1
A
A
A
B
B
C
D
D
D

期望输出为：

Col1 idx
A   1
A   2
A   3
B   1
B   2
C   1
D   1
D   2
D   3

在R中使用data.table可以通过df[, idx:=seq_len(.N), by=Col1]实现，但不清楚Python中的等价方法，仅了解numpy的linspace或arange函数，不知如何按分组操作。

解决方案 #

刚好我也常在pandas和data.table之间切换，这个需求用pandas其实有很简洁的实现，和你R里的写法逻辑完全对应：

方法1：groupby.cumcount()（最推荐）

这是最接近你data.table写法的等价方案，cumcount()会对每个分组内的行从0开始计数，加1就能得到从1起始的索引：

import pandas as pd

# 构造示例数据
df = pd.DataFrame({'Col1': ['A', 'A', 'A', 'B', 'B', 'C', 'D', 'D', 'D']})

# 添加分组内索引
df['idx'] = df.groupby('Col1').cumcount() + 1

print(df)

运行后就能得到你想要的结果，逻辑和R里的seq_len(.N)完全一致：cumcount()自动按Col1分组统计行数，加1是因为它默认从0开始计数，而你需要的是从1开始的序号。

方法2：自定义分组函数（更灵活）

如果需要更复杂的分组内逻辑（比如倒序编号、跳号等），可以用groupby.apply()手动生成序列：

df['idx'] = df.groupby('Col1').apply(lambda x: pd.Series(range(1, len(x)+1))).reset_index(drop=True)

这种方式代码稍长，但可以根据需求自由修改序列生成规则，比如把range(1, len(x)+1)改成range(len(x), 0, -1)就能实现倒序编号。

逻辑对应说明 #

• R data.table中的.N代表每个分组的行数，对应Python中分组后的len(x)或者x.shape[0]
• seq_len(.N)生成1到.N的序列，对应Python中的range(1, len(x)+1)，而cumcount()+1是pandas封装好的快捷方式

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