如何在Pandas中按ID提取column_1达指定值前后的column_2值

可以通过分组标记周期+周期内定位关键值的方式实现需求，核心思路是先按ID分组，为每个分组内的上升-回落阶段标记周期，再针对每个周期中column_1=5的行，提取对应的Before和After值。

实现代码 #

import pandas as pd

# 初始化数据
data = {'ID': ['A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'A', 'B', 'B', 'B', 'B', 'B', 'B', 'B', 'B', 'B', 'B'],
        'column_1': [0, 0, 0, 0, 0.1, 1, 1.5, 2, 3, 4, 4.5, 5, 4.9, 3, 2, 1.8, 1, 0, 0, 1, 3, 0, 1.3, 2, 3, 4.3, 4.8, 5, 4.2, 3.5, 3, 2.6, 2, 1.9, 1, 0, 0, 0, 0, 0, 0.1, 0.2, 0.3, 1, 2, 3, 5, 4, 2, 0.5, 0, 0],
        'column_2': [13,25,96,59,5,92,82,141,50,85,84,113,119,128,8,133,82,10,15,62,11,68,18,24,37,55,83,48,13,81,43,36,56,43,36,46,45,127,55,67,113,98,78,78,57,131,121,126,142,51,64,95]}

df = pd.DataFrame(data)

def process_group(g):
    # 初始化结果列
    g['Before'] = 0
    g['After'] = 0
    
    # 标记上升周期：从0转为非0时启动新周期，处理起始非0的情况
    start_signal = (g['column_1'] != 0) & (g['column_1'].shift(fill_value=0) == 0)
    g['cycle'] = start_signal.cumsum()
    
    # 遍历每个周期
    for cycle_num in g['cycle'].unique():
        cycle_data = g[g['cycle'] == cycle_num]
        if cycle_data.empty:
            continue
        
        # 定位当前周期中column_1=5的行
        target_indices = cycle_data[cycle_data['column_1'] == 5].index
        if target_indices.empty:
            continue
        
        # 计算Before值：周期从0开始则取第一个非0的column_2，否则取组首值
        if cycle_data.iloc[0]['column_1'] == 0:
            before_val = cycle_data[cycle_data['column_1'] != 0].iloc[0]['column_2']
        else:
            before_val = g.iloc[0]['column_2']
        
        # 计算After值：周期结束后第一个0的column_2
        last_non_zero_idx = cycle_data[cycle_data['column_1'] != 0].index[-1]
        after_segment = g.loc[last_non_zero_idx+1:]
        first_zero = after_segment[after_segment['column_1'] == 0]
        after_val = first_zero.iloc[0]['column_2'] if not first_zero.empty else 0
        
        # 赋值到目标行
        g.loc[target_indices, 'Before'] = before_val
        g.loc[target_indices, 'After'] = after_val
    
    # 清理临时列
    return g.drop('cycle', axis=1)

# 分组处理并输出结果
result = df.groupby('ID', group_keys=False).apply(process_group)
print(result)

关键步骤说明 #

1. 周期标记：通过shift()和cumsum()识别从0到非0的转换点，为每个上升-回落阶段分配唯一周期编号。
2. Before值提取：
   • 若周期起始为0，取周期内第一个非0行的column_2值；
   • 若周期起始非0（如ID B的首个周期），直接取分组的第一个column_2值。
3. After值提取：找到周期内最后一个非0行的位置，取之后第一个0行的column_2值。
4. 赋值填充：将计算得到的Before和After值填充到对应column_1=5的行，其余行保持0。

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