基于索引匹配的pandas双DataFrame乘法向量化优化实现

核心需求回顾 #

现有两个结构固定的数据集：

• df1：行、列均为两层MultiIndex，第一层为取值0_{N-1的id，第二层为取值0}T-1的时间周期，整体规模为(NT, NT)
• df2：普通二维权重矩阵，规模为(N,N)，每个位置(i,j)存储标量权重w_ij

需要计算的结果矩阵E规模为(N,N)，每个元素的计算规则为：
取出df1中所有行id为i、列id为j的(T*T)大小子矩阵，与对应权重w_ij逐元素相乘后累加求和。

原有实现的性能问题 #

当前采用双层嵌套循环，通过pandas的.loc接口按索引切片完成计算：

E = pd.DataFrame(index=range(N), columns=range(N))
# 低效循环逻辑
for i in range(N):
     for j in range(N):
        E.loc[i][j] = (df1.loc[i][j] * df2.loc[i][j]).to_numpy().sum()

该实现的主要开销来自两部分：

• 嵌套循环产生的Python层调度开销
• 每次.loc切片触发的多级索引匹配、pandas对象封装开销

在N=200、T=3的规模下，累计4万次循环已经产生明显耗时，无法支撑后续数百次重复计算的需求。

零循环向量化实现 #

整个计算完全可以通过numpy的连续内存操作完成，不需要任何Python层循环，速度可提升2个数量级以上，且结果和原逻辑完全等价。

计算逻辑拆解 #

从数学本质看，每个E[i,j]的计算可以拆为两步：

1. 先单独把df1中id块(i,j)的所有元素求总和
2. 将块总和与对应权重w_ij逐元素相乘

只要利用df1的多层索引连续排列的特性，完全可以一次性完成所有块的求和，无需逐一切片。

具体实现步骤 #

1. 直接提取df1的底层数值，转为连续存储的numpy二维数组，跳过pandas的索引封装：

   arr1 = df1.to_numpy()
   

2. 由于df1的行、列都是按「连续存储同一个id下所有时间维度元素」的规则排列，直接将形状为(N*T, N*T)的数组reshape为四维结构(N, T, N, T)，分别对应「行id、行内时间、列id、列内时间」四个维度：

   arr1_4d = arr1.reshape(N, T, N, T)
   

3. 对四维数组的「行内时间」「列内时间」两个维度直接求和，一次性得到所有id块的元素总和，输出形状为(N,N)：

   block_sum = arr1_4d.sum(axis=(1, 3))
   

4. 提取权重矩阵的numpy数组，与块总和逐元素相乘，得到最终结果：

   weight_mat = df2.to_numpy()
   E = block_sum * weight_mat
   

注意事项 #

如果df1的多层索引没有按id连续排序，需先对行、列做排序操作，保证同一个id的元素在内存中是连续块，避免reshape时把不同id的元素混排：

# 保证索引连续，仅当索引非有序时执行
df1 = df1.sort_index(axis=0).sort_index(axis=1)

效果对比 #

原嵌套循环在N=200、T=3的规模下单次计算耗时约数百毫秒，向量化实现下单次计算耗时可降低到微秒级，完全满足数百次重复计算的性能要求。

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