如何为DataFrame各分组按日期生成递增整数新列

原始数据 #

date          group         value
2022-11-01.     1              4
2022-11-02.     1              12
2022-11-03.     1              14
2022-11-04.     1              25
2021-11-01.     2              9
2021-11-02.     2              7
2019-10-01.     3              40
2022-10-02.     3              14

需求说明 #

需要新增一列new_col，在每个group分组内，按日期从小到大生成从0开始的递增整数，最终输出如下：

date          group         value      new_col
2022-11-01     1              4          0
2022-11-02     1              12         1
2022-11-03     1              14         2
2022-11-04     1              25         3
2021-11-01     2              9          0
2021-11-02     2              7          1
2019-10-01     3              40         0
2022-10-02     3              14         1

你尝试的问题代码 #

df.groupby('group')['date'].apply(lambda x: np.arange(0, len(x)+1))

这段代码存在三个问题：

1. len(x)+1生成的数组长度比分组内行数多1，无法匹配原数据行数
2. 未对分组内的日期进行排序，计数顺序可能和日期顺序不对应
3. apply返回的是分组级别的数组，直接赋值会导致结构不匹配，无法生成新列

正确实现方法 #

方法一：使用cumcount()（推荐） #

cumcount()是Pandas专门用于分组内计数的方法，从0开始递增，简洁高效。步骤如下：

1. 先将数据按group和date排序，确保分组内日期是递增顺序
2. 调用groupby + cumcount()生成新列

import pandas as pd
import numpy as np

# 处理日期列的格式（去掉末尾的点，转为日期类型）
df['date'] = pd.to_datetime(df['date'].str.rstrip('.'))

# 按分组和日期排序
df = df.sort_values(['group', 'date'])

# 生成new_col
df['new_col'] = df.groupby('group').cumcount()

方法二：手动用apply + np.arange #

如果要手动实现类似np.arange的逻辑，需要确保返回的数组长度和分组行数一致，并用explode()展开：

df['date'] = pd.to_datetime(df['date'].str.rstrip('.'))
df = df.sort_values(['group', 'date'])

df['new_col'] = df.groupby('group')['date'].apply(lambda x: np.arange(len(x))).explode()

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