如何一次性对多层索引Pandas DataFrame按组执行行除法？

问题背景 #

先构造一个多层索引的Pandas DataFrame：

import numpy as np
import pandas as pd

df1 = pd.DataFrame(np.random.rand(3, 3), index=['a', 'b', 'c'], columns=['col1', 'col2', 'col3'])
df2 = pd.DataFrame(np.random.rand(3, 3), index=['a', 'b', 'c'], columns=['col1', 'col2', 'col3'])
df3 = pd.DataFrame(np.random.rand(3, 3), index=['a', 'b', 'c'], columns=['col1', 'col2', 'col3'])
df = pd.concat([df1, df2, df3], axis=0, keys=['A', 'B', 'C'])

生成的DataFrame示例如下（你的随机数值会与示例不同）：

col1        col2        col3
A   a   0.893752    0.554021    0.492867
    b   0.319270    0.263366    0.542281
    c   0.082265    0.635637    0.796405
B   a   0.954748    0.684624    0.488293
    b   0.485414    0.966693    0.211348
    c   0.411648    0.989666    0.028412
C   a   0.701327    0.025172    0.320882
    b   0.073527    0.060885    0.111406
    c   0.169269    0.627686    0.438393

需求 #

• 将(A, b)和(A, c)行除以(A, a)行
• 将(B, b)和(B, c)行除以(B, a)行
• 将(C, b)和(C, c)行除以(C, a)行

个人尝试 #

我最初通过循环遍历顶层索引，分别做除法后合并结果：

idx = pd.IndexSlice
ratio_list = [df.loc[idx[x,:], :].div(df.loc[idx[x,'a'], :]) for x in ['A', 'B', 'C']]
ratio = pd.concat(ratio_list, axis=0)

生成的结果如下：

col1        col2        col3
A   a   1.000000    1.000000    1.000000
    b   0.357225    0.475371    1.100259
    c   0.092044    1.147315    1.615864
B   a   1.000000    1.000000    1.000000
    b   0.508422    1.412005    0.432830
    c   0.431159    1.445560    0.058186
C   a   1.000000    1.000000    1.000000
    b   0.104840    2.418784    0.347188
    c   0.241355    24.936324   1.366214

高效解决方案 #

以下是两种更优的实现方案：groupby-transform写法简洁易读，当DataFrame规模较大时，np.repeat方案的性能优势更显著。

方案一：Pandas groupby-transform（@Smordy提供） #

按顶层索引分组，通过transform('first')获取每组首行数据，再直接执行除法：

df.div(df.groupby(level=0).transform('first'))

方案二：Numpy repeat（@ouroboros1提供） #

提取每组首行数据，用np.repeat将其重复至与原DataFrame行数一致，再做除法：

df.div(np.repeat(df.loc[(slice(None), ['a']), :].to_numpy(), 3, axis=0))

性能测试 #

构造更大规模的DataFrame测试两种方案的性能：

nrow = 100  # 可调整行数测试
df = pd.DataFrame(np.random.rand(nrow*3, 3), 
                   columns=['col1', 'col2', 'col3'],
                   index=pd.MultiIndex.from_product([[*'ABC'], ['row' + str(ii) for ii in range(0, nrow)]]))

# Pandas groupby-transform 方案
%timeit df.div(df.groupby(level=0).transform('first'))

# Numpy repeat 方案
%timeit df.div(np.repeat(df.loc[(slice(None), ['row0']), :].to_numpy(), nrow, axis=0))

测试结果显示：当nrow=100时，np.repeat方案的运行速度约为groupby-transform的3倍；若增加DataFrame的列数，两者的速度差距还会进一步扩大。

内容的提问来源于stack exchange，提问作者data-monkey