使用Pandas计算分组内与最近条件行的差值

需求：为给定的DataFrame新增一列Diff，展示每行与所在group内最近（基于date）的indicator == True行的val差值，且indicator == True的行的Diff值为0。

示例输入DataFrame #

import pandas as pd

data = [['A', '2022-09-01', False, 2], ['A', '2022-09-02', False, 3], ['A', '2022-09-03', True, 1],
        ['A', '2022-09-05', False, 4], ['A', '2022-09-08', True, 4], ['A', '2022-09-09', False, 2],
        ['B', '2022-09-03', False, 4], ['B', '2022-09-05', True, 5], ['B', '2022-09-06', False, 7],
        ['B', '2022-09-09', True, 4], ['B', '2022-09-10', False, 2], ['B', '2022-09-11', False, 3]]
df = pd.DataFrame(data = data, columns = ['group', 'date', 'indicator', 'val'])

输入的DataFrame结构：

group        date  indicator  val
0      A  2022-09-01      False    2
1      A  2022-09-02      False    3
2      A  2022-09-03       True    1
3      A  2022-09-05      False    4
4      A  2022-09-08       True    4
5      A  2022-09-09      False    2
6      B  2022-09-03      False    4
7      B  2022-09-05       True    5
8      B  2022-09-06      False    7
9      B  2022-09-09       True    4
10     B  2022-09-10      False    2
11     B  2022-09-11      False    3

期望输出DataFrame #

data = [['A', '2022-09-01', False, 2, 1], ['A', '2022-09-02', False, 3, 2], ['A', '2022-09-03', True, 1, 0],
        ['A', '2022-09-05', False, 4, 3], ['A', '2022-09-08', True, 4, 0], ['A', '2022-09-09', False, 2, -2],
        ['B', '2022-09-03', False, 4, -1], ['B', '2022-09-05', True, 5, 0], ['B', '2022-09-06', False, 7, 2],
        ['B', '2022-09-09', True, 4, 0], ['B', '2022-09-10', False, 2, -2], ['B', '2022-09-11', False, 3, -1]]
df_desired = pd.DataFrame(data = data, columns = ['group', 'date', 'indicator', 'val', 'Diff'])

输出的DataFrame结构：

group        date  indicator  val  Diff
0      A  2022-09-01      False    2     1
1      A  2022-09-02      False    3     2
2      A  2022-09-03       True    1     0
3      A  2022-09-05      False    4     3
4      A  2022-09-08       True    4     0
5      A  2022-09-09      False    2    -2
6      B  2022-09-03      False    4    -1
7      B  2022-09-05       True    5     0
8      B  2022-09-06      False    7     2
9      B  2022-09-09       True    4     0
10     B  2022-09-10      False    2    -2
11     B  2022-09-11      False    3    -1

Diff列计算说明 #

以分组A为例：

• 第0行：2 - 1 = 1，对应最近的True行（第2行）
• 第1行：3 - 1 = 2，对应最近的True行（第2行）
• 第3行：4 - 1 = 3，对应最近的True行（第2行）
• 第5行：2 - 4 = -2，对应最近的True行（第4行）
• 所有indicator == True的行，Diff值固定为0，这些行作为差值计算的基准行

实现方案 #

通过以下步骤实现需求：

1. 将date列转换为datetime类型，确保日期排序逻辑正确
2. 按group分组，组内按date排序
3. 生成基准值列：仅保留indicator == True行的val，其余行填充为NaN；先向前填充最近的基准值，再向后填充（处理第一个True行之前的记录）
4. 计算val与填充后基准值的差值，最后将indicator == True行的Diff设为0

具体代码如下：

# 转换date列为datetime类型
df['date'] = pd.to_datetime(df['date'])

# 定义分组处理函数
def calculate_diff(group):
    # 组内按日期排序
    group = group.sort_values('date')
    # 创建基准值列：True行保留val，其余为NaN
    group['base_val'] = group['val'].where(group['indicator'])
    # 先向前填充最近基准值，再向后填充（处理组内首个True行之前的记录）
    group['base_val'] = group['base_val'].ffill().bfill()
    # 计算差值
    group['Diff'] = group['val'] - group['base_val']
    # 将True行的Diff设为0
    group.loc[group['indicator'], 'Diff'] = 0
    return group

# 应用分组函数并重置索引
df_result = df.groupby('group', group_keys=False).apply(calculate_diff).reset_index(drop=True)

print(df_result)

运行后即可得到符合要求的结果。

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