按MPRO和ID分组实现列乘法并满足目标值标记需求

需求明确 #

• 按MPRO和ID字段对数据分组
• 每组内计算Nuevo_I * Nuevo_P的乘积，规则如下：
  • 乘积大于8：保留计算结果
  • 乘积首次小于等于8：保留该结果；同组内该行之后的所有行，结果标记为'DONE'

输入数据示例 #

import pandas as pd

df = pd.DataFrame({
    'MPRO': ['A', 'A', 'A', 'B', 'B', 'B'],
    'ID': [1, 1, 1, 2, 2, 2],
    'Nuevo_I': [3, 2, 1, 4, 3, 2],
    'Nuevo_P': [3, 4, 3, 3, 2, 3]
})

期望输出 #

MPRO  ID  Nuevo_I  Nuevo_P Result
0    A   1        3        3      9
1    A   1        2        4      8
2    A   1        1        3   DONE
3    B   2        4        3     12
4    B   2        3        2      6
5    B   2        2        3   DONE

问题分析 #

常规的分组乘积计算无法处理「首次出现<=8后标记后续行」的逻辑，需要自定义分组处理函数来跟踪分组内的状态。

解决方案代码 #

def process_group(group):
    # 计算每组内的乘积
    products = group['Nuevo_I'] * group['Nuevo_P']
    # 找到分组中第一个乘积<=8的行索引
    first_leq8_pos = products[products <= 8].index.min()
    
    # 若分组内所有乘积都>8，直接返回乘积字符串
    if pd.isna(first_leq8_pos):
        return products.astype(str)
    
    # 否则，将第一个<=8之后的行标记为'DONE'
    result = products.astype(str)
    result.loc[first_leq8_pos + 1:] = 'DONE'
    return result

# 应用分组处理函数
df['Result'] = df.groupby(['MPRO', 'ID'], group_keys=False).apply(process_group)

代码说明 #

1. 自定义分组函数：process_group针对每个分组完成以下操作：
   • 计算当前组的乘积序列
   • 定位第一个满足<=8的行位置
   • 根据是否存在该位置，决定是返回全部乘积还是标记后续行为'DONE'
2. 分组应用：使用group_keys=False避免分组键混入结果，确保输出格式符合预期

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