加快 pandas 数据帧滚动加权平均值的方法

Question

我有一个很大的DataFrame，我需要在其上计算滚动行加权平均值。

我知道我可以执行以下操作：

import numpy as np
import pandas as pd

df = pd.DataFrame(np.random.rand(20000, 50))
weights = [1/9, 2/9, 1/3, 2/9, 1/9]  
rolling_mean = df.rolling(5, axis=1).apply(lambda seq: np.average(seq, weights=weights))

问题是在我的电脑上这大约需要 40 秒。有什么办法可以加快这个计算速度吗？

Answer 1

代码

创建一个新的数据帧，将 df 乘以权重 [0]，然后将 df 移位 1 并乘以权重 [1]，然后将 df 移位 2 并乘以权重 [2]，重复此过程，然后添加所有创建的数据帧数据帧放在一起，将加快该过程。

sum([df.shift(num, axis=1) * w for num, w in enumerate(weights)])

需要 0.05986 秒