基于同索引下另一DataFrame的值对DataFrame进行条件求和

问题场景 #

给定两个结构完全一致的DataFrame df1 和 df2，需要根据df1中元素等于指定值的条件，对df2对应位置的元素进行求和。

示例数据 #

先创建示例中的两个DataFrame：

import pandas as pd

df1 = pd.DataFrame(data=[[1,-1,5],[2,1,1],[3,0,0]], index=[0,1,2], columns=[0,1,2])
df2 = pd.DataFrame(data=[[1,1,3],[1,1,2],[0,2,1]], index=[0,1,2], columns=[0,1,2])

df1结构：

df2结构：

解决方案 #

方法1：布尔索引+双重求和

通过布尔筛选定位df1中符合条件的位置，再对df2对应元素求和：

# 计算df1元素为1时，df2对应元素的总和
sum_condition_1 = df2[df1 == 1].sum().sum()
print(sum_condition_1)  # 输出：4

# 计算df1元素为0时，df2对应元素的总和
sum_condition_0 = df2[df1 == 0].sum().sum()
print(sum_condition_0)  # 输出：3

说明：

• df1 == 1生成与原DataFrame同形状的布尔矩阵，标记符合条件的位置
• 用布尔矩阵索引df2时，不符合条件的位置会被设为NaN
• 两次sum()：第一次按列求和，第二次将列求和结果汇总，得到全局总和

方法2：使用where方法

where会将不符合条件的元素设为NaN，求和时自动忽略NaN：

sum_condition_1 = df2.where(df1 == 1).sum().sum()
sum_condition_0 = df2.where(df1 == 0).sum().sum()

方法3：Numpy向量化操作（大数据量推荐）

直接操作底层Numpy数组，执行效率更高：

import numpy as np

sum_condition_1 = np.sum(df2.values[df1.values == 1])
sum_condition_0 = np.sum(df2.values[df1.values == 0])

结果验证 #

• 条件为1时，df1中符合条件的位置对应df2的元素是1、1、2，总和为4
• 条件为0时，df1中符合条件的位置对应df2的元素是2、1，总和为3

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