基于索引值将指定行复制为新列（Pandas实现）

需求说明 #

现有一个包含Datetime、Val、Item列的DataFrame，需要实现：

• 筛选出Item == 'X'的行
• 对每一行，找到该行之后第一个Val值与当前行Val差值≥2的行索引
• 将找到的目标行的Datetime和Val值，分别存入当前行的Datetime_High和Val_High新列

原代码及问题 #

原代码尝试实现需求，但赋值逻辑错误，导致结果不符合预期：

import pandas as pd
import numpy as np

d = ({
    'Datetime' : ['2022-07-25 01:46:00+01:00','2022-07-25 01:47:00+01:00','2022-07-25 01:48:00+01:00','2022-07-25 01:49:00+01:00','2022-07-25 01:50:00+01:00','2022-07-25 01:51:00+01:00','2022-07-25 01:52:00+01:00','2022-07-25 01:53:00+01:00'],                 
    'Val' : [5,1,4,-2,8,4,3,10],   
    'Item' : ['X', np.NaN, np.NaN, np.NaN, 'X', np.NaN, 'X', np.NaN],                      
})

df = pd.DataFrame(data = d)

mask = df['Item'] == 'X'

next_high = np.where(mask, df.apply(lambda row: (df[df.index>=row.name].Val - row.Val)
                    .where(lambda v: v >= 2).first_valid_index(), axis=1), np.NaN)

# 原错误逻辑：将当前行的值赋值给目标行，与需求相反
df.loc[df.index[next_high], 'Val_High'] = df['Val']
df.loc[df.index[next_high], 'Datetime_High'] = df['Datetime']

错误分析 #

原代码的核心问题是赋值方向颠倒：

• 需求是把next_high对应行的Datetime/Val，放到Item='X'的当前行
• 但原代码是把当前行的Datetime/Val，放到next_high对应的行，完全不符合预期

修正后的代码 #

import pandas as pd
import numpy as np

d = ({
    'Datetime' : ['2022-07-25 01:46:00+01:00','2022-07-25 01:47:00+01:00','2022-07-25 01:48:00+01:00','2022-07-25 01:49:00+01:00','2022-07-25 01:50:00+01:00','2022-07-25 01:51:00+01:00','2022-07-25 01:52:00+01:00','2022-07-25 01:53:00+01:00'],                 
    'Val' : [5,1,4,-2,8,4,3,10],   
    'Item' : ['X', np.NaN, np.NaN, np.NaN, 'X', np.NaN, 'X', np.NaN],                      
})

df = pd.DataFrame(data = d)

mask = df['Item'] == 'X'

# 计算每个Item='X'行对应的后续第一个符合条件的索引
next_high = df.apply(lambda row: 
    (df.loc[row.name:, 'Val'] - row['Val']).ge(2).idxmax() 
    if mask[row.name] else np.nan, axis=1)

# 为Item='X'的行赋值：从next_high索引行取数据
df.loc[mask, 'Datetime_High'] = df.loc[next_high[mask], 'Datetime'].values
df.loc[mask, 'Val_High'] = df.loc[next_high[mask], 'Val'].values

# 查看结果
print(df)

预期输出 #

Datetime  Val Item              Datetime_High  Val_High
0  2022-07-25 01:46:00+01:00    5    X  2022-07-25 01:50:00+01:00       8.0
1  2022-07-25 01:47:00+01:00    1  NaN                        NaN       NaN
2  2022-07-25 01:48:00+01:00    4  NaN                        NaN       NaN
3  2022-07-25 01:49:00+01:00   -2  NaN                        NaN       NaN
4  2022-07-25 01:50:00+01:00    8    X  2022-07-25 01:53:00+01:00      10.0
5  2022-07-25 01:51:00+01:00    4  NaN                        NaN       NaN
6  2022-07-25 01:52:00+01:00    3    X  2022-07-25 01:53:00+01:00      10.0
7  2022-07-25 01:53:00+01:00   10  NaN                        NaN       NaN

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