如何基于分组与列条件在Pandas中新增test列

原始数据 #

首先，你的原始数据生成代码如下：

import pandas as pd
import numpy as np

data = {
    'id': [1, 2, 3, 4, 5, 6, 7],
    'date': ['2019-02-01', '2019-02-10', '2019-02-25', '2019-03-05', '2019-03-16', '2019-04-05', '2019-05-15'],
    'date_difference': [None, 9, 15, 11, 10, 19, 40],
    'number': [1, 0, 1, 0, 0, 0, 0],
    'text': ['A', 'A', 'A', 'A', 'A', 'B', 'B']
}

df = pd.DataFrame(data)
# 先将date转为datetime类型，方便排序
df['date'] = pd.to_datetime(df['date'])

对应的初始表格：

问题分析 #

你需要生成test列，规则为：

• 按text分组，每组内按日期降序处理
• 步长初始为1，遇到number == 1时步长加1；无1则步长保持1
• 你的原有代码错误地将text和number一起分组，且未实现步长动态累加的逻辑，因此无法得到正确结果

正确实现代码 #

通过分组后自定义函数处理，实现步长的动态更新：

def compute_test(group):
    # 分组内按日期降序排序
    sorted_group = group.sort_values('date', ascending=False).copy()
    step = 1
    test_values = []
    # 遍历排序后的number列，先赋值当前步长，再判断是否累加步长
    for num in sorted_group['number']:
        test_values.append(step)
        if num == 1:
            step += 1
    sorted_group['test'] = test_values
    # 恢复原数据的索引顺序
    return sorted_group.loc[group.index]

# 应用分组处理
df = df.groupby('text', group_keys=False).apply(compute_test)

最终结果 #

运行后得到的表格如下：

逻辑说明 #

1. 每组按日期降序排序后，从最新日期开始遍历
2. 先将当前步长赋值给test，再判断当前行number是否为1，若为1则步长加1
3. 最后恢复原数据的索引顺序，保证结果与原数据行对应

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