按条件映射计算Pandas DataFrame中指定单元格的数值总和

问题说明 #

现有如下Pandas DataFrame，需按以下两个条件分别计算对应单元格的数值总和：

• 条件1：LOC取值为8、21，且仅计算TYPE为x、y、z的列，预期结果为27
• 条件2：LOC取值为30、16，且仅计算TYPE为x、y、z、q的列，预期结果为21

示例DataFrame构造代码：

import pandas as pd

data = {'LOC': [5, 8, 16, 21, 30, 35],
        'x': [5, 2, 3, 5, 2, 2],
        'y': [6, 3, 1, 6, 2, 1],
        'z': [7, 3, 2, 8, 3, 4],
        'q': [8, 4, 2, 9, 6, 5],
        'p': [10, 5, 3, 1, 7, 3]
       }
data_frame = pd.DataFrame(data)

解决方案 #

方法一：针对单个条件直接计算 #

条件1计算

先筛选LOC为8或21的行，再选取x、y、z列，最后求和：

# 条件1总和
cond1_sum = data_frame[data_frame['LOC'].isin([8, 21])][['x', 'y', 'z']].sum().sum()
print(cond1_sum)  # 输出27

条件2计算

同理，筛选LOC为30或16的行，选取x、y、z、q列求和：

# 条件2总和
cond2_sum = data_frame[data_frame['LOC'].isin([30, 16])][['x', 'y', 'z', 'q']].sum().sum()
print(cond2_sum)  # 输出21

方法二：封装通用函数（复用性更强） #

如果需要频繁处理类似条件，可以写一个通用函数，传入LOC列表和TYPE列列表，直接返回总和：

def calculate_sum(df, loc_values, type_columns):
    # 筛选LOC符合条件的行，选取指定TYPE列，计算嵌套总和
    return df[df['LOC'].isin(loc_values)][type_columns].sum().sum()

# 调用函数计算条件1
cond1_sum = calculate_sum(data_frame, [8, 21], ['x', 'y', 'z'])
# 调用函数计算条件2
cond2_sum = calculate_sum(data_frame, [30, 16], ['x', 'y', 'z', 'q'])

print(cond1_sum)  # 27
print(cond2_sum)  # 21

代码说明 #

• df['LOC'].isin(loc_values)：筛选LOC列值在指定列表中的行
• [type_columns]：选取需要计算的TYPE列（即x、y等列）
• 第一个sum()是对每列求和，第二个sum()是把各列的和再相加，得到所有符合条件的单元格的总和

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