基于pandas DataFrame其他列值在s_w列查找1的模式并计算新列X的值

• 测试用DataFrame结构如下：

import pandas as pd
df = pd.DataFrame( {
    "p":[1,1,1,1,1,1,1,1,1,1,1,1,1],
    "l" :[1,1,1,1,1,1,1,1,1,1,1,1,1],
    "w":[1,2,3,4,5,6,7,8,9,10,11,12,12],
    "s_w":[1,1,0,0,0,0,1,0,0,0,0,1,1],
    "key" :[1,1,0,0,0,1,0,1,0,0,0,0,1],
    "v":[2,2,5,3,4,5,5,1,2,3,4,5,4]
})

• 处理规则：
  • 仅处理key==1的行，其余行新增列X直接设为NaN
  • 对符合条件的行，分别在前序行、后序行中查找s_w==1的记录，两边各需要至少2条匹配记录，任意一边不满足则X设为NaN
  • 匹配到的4条记录的v值求和即为对应行的X值，s_w的1无需连续，允许中间存在间隔

# 提前提取所有s_w=1的行的索引和对应v值，避免重复计算
sw1_list = df.loc[df['s_w'] == 1, ['v']].reset_index().values.tolist()
sw1_index = [x[0] for x in sw1_list]
sw1_v = [x[1] for x in sw1_list]

def get_x(row):
    if row['key'] != 1:
        return pd.NA
    cur_id = row.name
    # 取前序最近2个s_w=1的v值
    prev_match = [sw1_v[i] for i, idx in enumerate(sw1_index) if idx < cur_id]
    if len(prev_match) < 2:
        return pd.NA
    # 取后序最早2个s_w=1的v值
    next_match = [sw1_v[i] for i, idx in enumerate(sw1_index) if idx > cur_id]
    if len(next_match) < 2:
        return pd.NA
    return sum(prev_match[-2:] + next_match[:2])

df['X'] = df.apply(get_x, axis=1)

最终生成的X列值对应如下：

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