基于Pandas按条件实现DataFrame中TRGQTY列值乘法运算

需要在Pandas DataFrame中实现：选取SUBPACKNR列非空行的TRGQTY值，与所有BALICI_INSTRUKCE列匹配对应层级+1（例如PV2对应PV3，即原层级为PVn则匹配PV(n+1)）的行的TRGQTY值相乘。

示例DataFrame #

import pandas as pd

data = {
    'BALICI_INSTRUKCE': ['PV2', 'PV3', 'PV2', 'PV3', 'PV1'],
    'SUBPACKNR': ['SP1', '', 'SP2', '', ''],
    'TRGQTY': [2, 3, 4, 5, 10]
}
df = pd.DataFrame(data)

输出的初始DataFrame：

BALICI_INSTRUKCE SUBPACKNR  TRGQTY
0              PV2       SP1       2
1              PV3                3
2              PV2       SP2       4
3              PV3                5
4              PV1               10

错误代码（用户提供） #

# 错误点：错误匹配层级逻辑，且未正确累加对应层级的所有TRGQTY值
df['RESULT'] = df.apply(
    lambda row: row['TRGQTY'] * df[df['BALICI_INSTRUKCE'] == f"PV{int(row['BALICI_INSTRUKCE'][2:])+1}"]['TRGQTY'].sum()
    if pd.notna(row['SUBPACKNR']) and row['SUBPACKNR'] != ''
    else row['TRGQTY'],
    axis=1
)

当前错误结果 #

BALICI_INSTRUKCE SUBPACKNR  TRGQTY  RESULT
0              PV2       SP1       2       8  # 预期应为 2*(3+5)=16
1              PV3                3       3
2              PV2       SP2       4      16  # 预期应为 4*(3+5)=32
3              PV3                5       5
4              PV1               10      10

期望输出 #

BALICI_INSTRUKCE SUBPACKNR  TRGQTY  RESULT
0              PV2       SP1       2      16  # 2*(3+5)
1              PV3                3       3
2              PV2       SP2       4      32  # 4*(3+5)
3              PV3                5       5
4              PV1               10      10

正确实现方案 #

步骤1：提取层级数字并定义目标层级 #

先从BALICI_INSTRUKCE中提取层级数字，计算需要匹配的+1层级：

# 提取PV后的数字并转为整数
df['LEVEL'] = df['BALICI_INSTRUKCE'].str.extract(r'PV(\d+)').astype(int)
df['TARGET_LEVEL'] = df['LEVEL'] + 1

步骤2：预计算各层级的TRGQTY总和 #

提前按层级分组求和，避免在apply中重复查询DataFrame，提升性能：

# 按BALICI_INSTRUKCE分组，计算每组TRGQTY的总和
level_trg_sum = df.groupby('BALICI_INSTRUKCE')['TRGQTY'].sum()

# 构建原层级到目标层级总和的映射
target_sum_map = {
    f"PV{level}": level_trg_sum.get(f"PV{level+1}", 0)
    for level in df['LEVEL'].unique()
}

步骤3：应用乘法逻辑 #

对SUBPACKNR非空的行执行乘法，其他行保持原TRGQTY值：

df['RESULT'] = df.apply(
    lambda row: row['TRGQTY'] * target_sum_map.get(row['BALICI_INSTRUKCE'], 1)
    # 处理SUBPACKNR为空字符串或空格的情况
    if pd.notna(row['SUBPACKNR']) and row['SUBPACKNR'].strip() != ''
    else row['TRGQTY'],
    axis=1
)

# 可选：删除中间辅助列
df = df.drop(['LEVEL', 'TARGET_LEVEL'], axis=1)

完整验证代码 #

import pandas as pd

data = {
    'BALICI_INSTRUKCE': ['PV2', 'PV3', 'PV2', 'PV3', 'PV1'],
    'SUBPACKNR': ['SP1', '', 'SP2', '', ''],
    'TRGQTY': [2, 3, 4, 5, 10]
}
df = pd.DataFrame(data)

# 提取层级数字
df['LEVEL'] = df['BALICI_INSTRUKCE'].str.extract(r'PV(\d+)').astype(int)
df['TARGET_LEVEL'] = df['LEVEL'] + 1

# 预计算各层级TRGQTY总和
level_trg_sum = df.groupby('BALICI_INSTRUKCE')['TRGQTY'].sum()
target_sum_map = {
    f"PV{level}": level_trg_sum.get(f"PV{level+1}", 0)
    for level in df['LEVEL'].unique()
}

# 计算结果
df['RESULT'] = df.apply(
    lambda row: row['TRGQTY'] * target_sum_map.get(row['BALICI_INSTRUKCE'], 1)
    if pd.notna(row['SUBPACKNR']) and row['SUBPACKNR'].strip() != ''
    else row['TRGQTY'],
    axis=1
)

# 删除辅助列
df = df.drop(['LEVEL', 'TARGET_LEVEL'], axis=1)

print(df)

输出结果与期望完全一致：

BALICI_INSTRUKCE SUBPACKNR  TRGQTY  RESULT
0              PV2       SP1       2      16
1              PV3                3       3
2              PV2       SP2       4      32
3              PV3                5       5
4              PV1               10      10

关键说明 #

• 预计算层级总和避免了apply中的重复查询，在大数据集下性能提升明显
• 用strip()处理SUBPACKNR的空格问题，避免误判非空值
• 若某层级无对应+1层级（如示例中的PV1），默认乘以0，可通过修改target_sum_map中的get默认值调整逻辑

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