带双重条件的Pandas DataFrame求和方案咨询

原始数据 #

创建数据的代码：

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

对应数据：

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

需求说明 #

对每一行的唯一(i,k)对，计算满足以下条件的X_ik之和：

• 条件1：其他行的k等于当前行的i
• 条件2：其他行的i不等于当前行的k
  若无符合条件的数值，求和结果为0。

示例参考：

(i,k) sum              explanation
(1,2): 0               sum all X_i1, but X_21 (if it exists)    
(1,3): 1               sum all X_i1, but X_31 (if it exists)   
(2,1): 0               sum all X_i2, but X_12 (if it exists)
(2,3): 4               sum all X_i2, but X_32 (if it exists)
(3,4): 0+14            sum all X_i3, but X_43 (if it exists) 
(3,5): 0+14            sum all X_i3, but X_53 (if it exists)
(3,6): 0+14            sum all X_i3, but X_63 (if it exists)

当前实现的问题 #

现有代码仅实现了条件1的求和，未排除条件2的情况，导致部分行结果错误：

s = df.groupby('k')['X_ik'].sum()
df['sum'] = df['i'].map(s)
display(df)

错误结果：

i   k   X_ik sum
0   1   2   4    1
1   1   3   0    1
2   2   1   1    4
3   2   3   14   4
4   3   4   2    14
5   3   5   12   14
6   3   6   5    14

向量化解决方案 #

方案1：字典映射+apply（高效易用） #

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

# 计算每个k对应的X_ik总和（满足条件1的总和）
sum_by_k = df.groupby('k')['X_ik'].sum()

# 创建(i,k)到X_ik的快速映射字典
ik_mapping = df.set_index(['i', 'k'])['X_ik'].to_dict()

# 计算最终结果：总和减去需要排除的X_ik（如果存在）
df['sum'] = df.apply(
    lambda row: sum_by_k.get(row['i'], 0) - ik_mapping.get((row['k'], row['i']), 0),
    axis=1
)

display(df)

得到正确结果：

i   k   X_ik  sum
0   1   2   4     0
1   1   3   0     1
2   2   1   1     0
3   2   3   14    4
4   3   4   2     14
5   3   5   12    14
6   3   6   5     14

方案2：全向量化合并（适合超大数据量） #

完全避免逐行操作，通过表合并实现：

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

# 步骤1：计算条件1的总和
sum_by_k = df.groupby('k')['X_ik'].sum()
df['total'] = df['i'].map(sum_by_k).fillna(0)

# 步骤2：创建反向映射表，存储(k,i)对应的X_ik
reverse_df = df.rename(columns={'i':'k', 'k':'i'})[['i','k','X_ik']].rename(columns={'X_ik':'exclude_val'})

# 步骤3：合并表，找到需要排除的值
df = df.merge(reverse_df, on=['i','k'], how='left')

# 步骤4：计算最终结果并清理临时列
df['sum'] = df['total'] - df['exclude_val'].fillna(0)
df = df.drop(columns=['total', 'exclude_val'])

display(df)

该方案利用Pandas的向量化合并操作，性能更优，适合处理百万级以上数据量。

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