如何在DataFrame中按组行值对应相减并新增列（无需拆分数据）

现有如下结构的DataFrame：

stage   val A       temp    shape
stage_0 10  red     hot     square  
stage_1 11  red     hot     square  
stage_2 12  red     hot     square  
stage_0 10  red     temp1   square  
stage_1 11  red     temp1   square  
stage_2 12  red     temp1   square  
stage_0 10  red     temp2   square  
stage_1 11  red     temp2   square  
stage_2 12  red     temp2   square  
stage_0 10  blue    temp1   square  
stage_1 11  blue    temp1   square  
stage_2 12  green   temp1   square  
stage_0 10  blue    temp2   square  
stage_1 11  blue    temp2   square  
stage_2 12  green   temp2   square  

核心信息：

• 列A有3个唯一值，temp有4个，shape有3个
• 每个A+temp+shape的唯一组合对应3个不同stage的val值，DataFrame共108行

需求：
将满足(A == 'red') & (temp == 'hot') & (shape == 'square')的行组作为基准，按stage一一对应，计算其他同temp、同shape行组的val与基准值的差值，结果存入新列out；需覆盖所有A与shape='square'的组合，且不创建新DataFrame完成操作。

预期结果示例：

stage   val A       temp    shape    out
stage_0 10  red     hot     square   0  
stage_1 11  red     hot     square   0
stage_2 12  red     hot     square   0
stage_0 11  red     temp1   square   1
stage_1 12  red     temp1   square   1
stage_2 13  red     temp1   square   1
stage_0 14  red     temp2   square   4
stage_1 15  red     temp2   square   4
stage_2 16  red     temp2   square   4
stage_0 10  blue    temp1   square   10
stage_1 11  blue    temp1   square   11
stage_2 12  green   temp1   square   12
stage_0 10  blue    temp2   square   10
stage_1 11  blue    temp2   square   11
stage_2 12  green   temp2   square   12

可以通过映射匹配结合原地赋值实现需求，下面提供两种高效实现方式：

方式一：基于map的高效匹配 #

# 提取基准组的val值，构建stage到val的映射字典
base_map = df[(df['A'] == 'red') & (df['temp'] == 'hot') & (df['shape'] == 'square')].set_index('stage')['val']

# 筛选shape=square的行，匹配对应stage的基准值并计算差值
mask = df['shape'] == 'square'
df.loc[mask, 'out'] = df.loc[mask, 'val'] - df.loc[mask, 'stage'].map(base_map)

# 给非shape=square的行填充默认值（按需调整，比如NaN或0）
df['out'] = df['out'].fillna(0)

方式二：逐行apply实现（适合小数据集） #

# 提取基准组的val值，按stage索引
base_vals = df[(df['A'] == 'red') & (df['temp'] == 'hot') & (df['shape'] == 'square')].set_index('stage')['val']

# 针对shape=square的行，逐行计算差值并赋值
df.loc[df['shape'] == 'square', 'out'] = df.loc[df['shape'] == 'square', :].apply(
    lambda row: row['val'] - base_vals[row['stage']], axis=1
)

# 非目标行填充默认值
df['out'] = df['out'].fillna(0)

核心逻辑说明 #

1. 基准值提取：先筛选出作为计算基准的行组，将stage设为索引，得到每个stage对应的基准val，方便后续快速匹配。
2. 原地赋值：仅针对shape='square'的行进行计算，直接通过df.loc赋值给新列out，全程在原DataFrame上修改，不生成新对象。
3. 补全默认值：根据需求给非目标行的out列填充默认值，避免出现缺失值。

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