基于分组DataFrame长度更新newtime列的Pandas实现方案问询

问题背景 #

原始数据 #

生成该数据的Pandas代码：

import pandas as pd
df = pd.DataFrame({
    'ID': ['aaa', 'aaa', 'aaa', 'bbb', 'ccc', 'ccc'],
    'closed': ['22:00', '22:00', '22:00', '23:00', '22:00', '22:00'],
    'set': ['19:00', '19:30', '20:00', '22:00', '19:00', '19:30'],
    'message_time': ['19:05', '19:40', '21:00', '22:10', '19:05', '19:40']
})

需求 #

新增newtime列，计算规则：

• 按ID分组，若分组长度为1：newtime = closed - set（差值以分钟为单位，格式为"X min"）
• 若分组长度大于1：
  • 非最后一行：newtime = set[row+1] - set[row]
  • 最后一行：newtime = closed[row] - set[row]

预期输出 #

最优实现代码 #

import pandas as pd

# 1. 将时间字符串转为datetime类型，便于计算差值
df['closed'] = pd.to_datetime(df['closed'], format='%H:%M')
df['set'] = pd.to_datetime(df['set'], format='%H:%M')

# 2. 生成分组辅助信息
grouped = df.groupby('ID')
# 标记每个分组的长度
df['group_size'] = grouped['ID'].transform('size')
# 标记当前行是否为分组内最后一行
df['is_last_row'] = grouped.cumcount(ascending=False) == 0
# 获取分组内下一行的set值（最后一行会得到NaN）
df['next_set'] = grouped['set'].shift(-1)

# 3. 计算newtime
def compute_newtime(row):
    if row['group_size'] == 1:
        time_diff = row['closed'] - row['set']
    else:
        if row['is_last_row']:
            time_diff = row['closed'] - row['set']
        else:
            time_diff = row['next_set'] - row['set']
    # 转换为分钟数并格式化字符串
    return f"{int(time_diff.total_seconds() // 60)} min"

df['newtime'] = df.apply(compute_newtime, axis=1)

# 4. 清理临时辅助列（可选，若不需要保留）
df.drop(['group_size', 'is_last_row', 'next_set'], axis=1, inplace=True)

# 5. 还原时间列为原始字符串格式（可选）
df['closed'] = df['closed'].dt.strftime('%H:%M')
df['set'] = df['set'].dt.strftime('%H:%M')

# 输出结果
print(df)

代码解释 #

1. 时间类型转换：把closed和set列转为datetime是计算时间差的前提，Pandas的时间类型支持直接做减法运算。
2. 分组辅助列：
   • group_size：快速判断分组是否只有一行
   • is_last_row：通过cumcount(ascending=False)倒序计数，等于0的就是最后一行
   • next_set：用shift(-1)获取分组内下一行的set值，非最后一行可直接用于差值计算
3. 差值计算：通过apply逐行匹配规则，计算时间差后转为分钟数，再格式化为要求的字符串。
4. 可选步骤：清理临时列和还原时间格式，根据实际需求选择是否执行。

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