Pandas分组下基于前后非连续符合条件行计算新列值的实现咨询

首先导入依赖并构造测试数据：

import pandas as pd
import numpy as np

# 构造带check列的原始DataFrame
df = pd.DataFrame({
    'p': [1,1,1,1,1,2,2,2,2,2,2],
    'l': [1,1,1,1,1,1,1,1,1,1,1],
    'w': [1,2,3,4,5,1,2,3,4,5,6],
    's_w': [1,1,0,1,1,1,0,0,0,1,1],
    'check': [0,0,0,0,1,0,0,0,0,0,0],
    'v': [2,2,5,5,5,1,2,3,4,5,4]
})

核心实现逻辑 #

自定义分组处理函数，完全匹配需求规则：

def process_group(group):
    # 提取当前分组内所有s_w=1的行，按w升序排序保证顺序正确
    sw1_rows = group.query("s_w == 1").sort_values("w").reset_index()
    total_sw1 = len(sw1_rows)
    # 初始化c_s列为空值
    group["c_s"] = np.nan
    
    for idx, row in group.iterrows():
        # 规则1：仅check=1的行参与计算
        if row["check"] != 1:
            continue
        # 规则2：当前行s_w必须为1
        if row["s_w"] != 1:
            continue
        # 查找当前行在s_w=1序列中的位置
        current_pos_series = sw1_rows[sw1_rows["index"] == idx].index
        if current_pos_series.empty:
            continue
        current_pos = current_pos_series[0]
        # 检查前后是否至少各有2个s_w=1的行
        if current_pos < 2 or (total_sw1 - current_pos - 1) < 2:
            continue
        # 计算前2个+后2个s_w=1行的v值之和
        sum_val = (
            sw1_rows.loc[current_pos - 2, "v"] 
            + sw1_rows.loc[current_pos - 1, "v"] 
            + sw1_rows.loc[current_pos + 1, "v"] 
            + sw1_rows.loc[current_pos + 2, "v"]
        )
        group.loc[idx, "c_s"] = sum_val
    return group

# 按p、l分组应用处理函数
df = df.groupby(["p", "l"], group_keys=False).apply(process_group)

结果验证 #

执行完成后打印print(df)即可得到计算后的结果，如果需要匹配你给出的期望输出，只需要调整原始数据中check列标记为1的行位置即可。

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