基于其他列条件对Pandas DataFrame列求和的优化实现

原数据集 #

import pandas as pd
df = pd.DataFrame({"i": [1,1,2,2,3,3],
                   "k": [2,3,1,3,4,5],
                   "Y_ik": [0,0,1,14,2,12]})
display(df)

输出：

i   k   Y_ik
0   1   2   0
1   1   3   0
2   2   1   1
3   2   3   14
4   3   4   2
5   3   5   12

需求说明 #

对每行计算满足以下条件的Y_ik求和值，新增列new_column存储结果：

• 基于(i,k)唯一分组的前提
• 求和范围为当前行的k值未出现在其他行的(i,k)组合中的所有Y_ik

预期结果：

i   k   Y_ik new_column
0   1   2   0    28
1   1   3   0    14
2   2   1   1    28
3   2   3   14   15
4   3   4   2    13
5   3   5   12   3

求和逻辑示例 #

(i,k)
(1,2): 0+14+2+12 = 28               
(1,3): 0+2+12 = 14          
(2,1): 0+14+2+12 = 28
(2,3): 1+2+12 = 15
(3,4): 0+1+12 = 13
(3,5): 0+1+2 = 3

要求方法高效，适用于大数据集，避免for循环等低效方式。

尝试过的代码及问题 #

尝试使用分组apply，但结果仅返回原Y_ik值：

df.groupby(["i","k"]).apply(lambda x : (x[~x["k"].isin(df[["i","k"]])]["Y_ik"].sum()))

输出：

i  k
1  2     0
   3     0
2  1     1
   3    14
3  4     2
   5    12

扩展场景的问题 #

原解决方案（@mozway提出）在扩展数据集下会产生NaN值：

扩展数据集 #

import pandas as pd
df = pd.DataFrame({"i": [1,1,2,2,3,3,8],
                   "k": [2,3,1,3,4,5,15],
                   "Y_ik": [0,0,1,14,2,12,1]})
display(df)

输出：

i   k   Y_ik
0   1   2   0
1   1   3   0
2   2   1   1
3   2   3   14
4   3   4   2
5   3   5   12
6   8   15  1

使用原方案代码 #

s = df.groupby('k')['Y_ik'].sum()
df['new_column'] = df['i'].map(s).add(df['k'].map(s)).rsub(s.sum())
display(df)

问题结果 #

新增行出现NaN：

i   k   Y_ik    new_column
0   1   2   0       29.0
1   1   3   0       15.0
2   2   1   1       29.0
3   2   3   14      16.0
4   3   4   2       14.0
5   3   5   12      4.0
6   8   15  1       NaN

解决方案 #

要解决NaN问题，需处理i值未在k列出现时的映射空值，修改后的高效代码如下：

# 计算每个k对应的Y_ik总和
k_sum = df.groupby('k')['Y_ik'].sum()
# 计算所有Y_ik的总和
total_sum = df['Y_ik'].sum()
# 映射i和k对应的总和，空值填充为0，再计算new_column
df['new_column'] = total_sum - df['i'].map(k_sum).fillna(0) - df['k'].map(k_sum).fillna(0)

验证扩展数据集结果 #

运行后得到正确结果：

i   k   Y_ik  new_column
0   1   2   0          29.0
1   1   3   0          15.0
2   2   1   1          29.0
3   2   3   14         16.0
4   3   4   2          14.0
5   3   5   12          4.0
6   8   15  1          28.0

逻辑说明 #

核心思路是利用总和减去排除项实现高效计算：

1. 先计算所有Y_ik的总和total_sum
2. 排除项包括：所有k等于当前行i的Y_ik总和（这些行的k匹配当前行的i，属于需要排除的情况），加上当前行k对应的Y_ik总和
3. 用fillna(0)处理i未在k列出现的场景，避免产生NaN

该方法完全基于Pandas向量化操作，无循环，效率极高，适配大数据集。

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