如何为DataFrame按ID分组生成sub列：用INT减对应time0行的LMP

需求说明 #

在已生成FM列的DataFrame中，创建名为sub的新列，计算规则为：

• 对每个唯一ID，找到其FM列值为time0的行对应的LMP值
• 用每行的INT列值减去该ID对应的上述LMP值，得到sub列的值

初始数据与已完成代码 #

import pandas as pd
import numpy as np

data = {
    'ID': [0, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2],
    'VIS': [0.0, 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0],
    'STA': [float('NaN'), 4.0, 7.0, 7.0, 7.0, 7.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0],
    'LMP': [float('NaN'), -35.0, 411.0, 773.0, 1143.0, 1506.0, float('NaN'), float('NaN'), float('NaN'), float('NaN'), float('NaN'), float('NaN')],
    'INT': [0.0, 0.0, 413.0, 777.0, 1171.0, 1509.0, 1967.0, 2310.0, 2627.0, 2970.0, 3357.0, 3768.0],
    'FM': [-1, -1, "time0", -1, -1, "time0", -1, -1, -1, -1, -1,-1]
}

sorted_data = pd.DataFrame(data)

# 已完成的FM列生成代码
sorted_data['FM'] = np.nan
for id in sorted_data['ID'].unique():
    filter_condition = (sorted_data['ID'] == id) & (~sorted_data['LMP'].isnull())
    if filter_condition.any():
        last_row_index = sorted_data.loc[filter_condition].index[-1]
        sorted_data.loc[last_row_index, 'FM'] = 'time0'

sorted_data['FM'] = sorted_data['FM'].fillna(-1)

实现'sub'列的代码 #

步骤1：构建ID与对应time0的LMP值映射 #

先提取每个ID下FM='time0'的LMP值，存为字典：

# 提取各ID对应的time0行的LMP值
time0_lmp_map = sorted_data[sorted_data['FM'] == 'time0'].set_index('ID')['LMP'].to_dict()

步骤2：计算sub列 #

用ID匹配上述映射中的LMP值，再用INT减去该值：

# 计算sub列
sorted_data['sub'] = sorted_data['INT'] - sorted_data['ID'].map(time0_lmp_map)

验证结果 #

生成的sub列计算结果如下（注：预期输出中ID=2对应的LMP值存在笔误，实际以数据中time0行的1506.0为准）：

• 行0：NaN（无对应time0的LMP值）
• 行1：0.0 - 411.0 = -411.0
• 行2：413.0 - 411.0 = 2.0
• 行3：777.0 - 1506.0 = -729.0
• 行4：1171.0 - 1506.0 = -335.0
• 行5：1509.0 - 1506.0 = 3.0
• 行6：1967.0 - 1506.0 = 461.0
• 行7：2310.0 - 1506.0 = 804.0
• 行8：2627.0 - 1506.0 = 1121.0
• 行9：2970.0 - 1506.0 = 1464.0
• 行10：3357.0 - 1506.0 = 1851.0
• 行11：3768.0 - 1506.0 = 2262.0

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