基于滚动偏移计算加权平均值的更快方法

问题描述

我有一个像 df = pd.DataFrame({'price': [100,101,99,95,97,88],'qty': [12,5,1,3,3]}) 这样的示例 df。我想计算 (price * qty / qty) 的滚动 5 数量平均值，并且所需的输出是 100,100.6,96.2,91.2。

不幸的是，我目前没有一个很好的方法来计算这个，我有一个很慢的方法来计算数量的累积总和，然后 df.qty_cumsum[(df.qty_cumsum<= x.qty_cumsum- 5)].argmax() 返回数量的最大 arg - 5，然后我就可以用它来计算第二步的加权平均值。

谢谢

解决方法

一种选择是重复 price，然后将 rolling 与行和 groupby 索引，取 last：

np.repeat(df['price'],df['qty']).rolling(5).mean().groupby(level=0).last()

输出：

0    100.0
1    101.0
2    100.6
3     97.0
4     96.2
5     91.2
Name: price,dtype: float64

附言如果您有较大的 qty 值，通过将 qty 剪裁为 5 来提高效率也可能是有意义的（例如，如果它是 5 或 12，则没有区别）：>

np.repeat(df['price'],np.clip(df['qty'],5)
         ).rolling(5).mean().groupby(level=0).last()