python - Pandas 数据框在没有 for 循环的情况下迭代行

Question

我有一个大约 300 万行的数据框，如下所示:

   date        size  price
0  2018-08-01  100   220
1  2018-08-01  110   245
2  2018-08-01  125   250
3  2018-08-02  110   210
4  2018-08-02  120   230
5  2018-08-02  150   260
6  2018-08-03  115   200

对于每一行，它都是一个项目的交易。我们有交易日期、商品尺寸和价格。

现在我想添加一个名为 avg_price 的列，这样一笔交易/行的 avg_price 就是最后一天的 k 笔交易的平均值最接近的尺寸(与 k 最近邻的想法非常相似)。

例如，当 k = 2 时，上面最后一行的 avg_price 应为 (210+230)/2=220，因为最近的 2 笔交易的规模为 110 和 120，对应的价格为 210 和 230。

预期输出应该是:

   date        size  price avg_price
0  2018-08-01  100   220   NA
1  2018-08-01  110   245   NA
2  2018-08-01  125   250   NA
3  2018-08-02  110   210   (220+245)/2
4  2018-08-02  120   230   (245+250)/2
5  2018-08-02  150   260   (245+250)/2
6  2018-08-03  115   200   (210+230)/2

我编写了一个 for 循环来迭代每一行，首先挑选出最后一天的所有交易，然后按大小差异排序并计算前 k 项的平均值。然而，正如预期的那样，这非常慢。有人能指出一种更“矢量化”的方法吗？谢谢。

更新:每天的交易数量不固定，大约在 300 笔左右。

Answer 1

我将 dfa 称为原始数据帧。首先在 dfb 中创建稍后需要的数据 merge_asof

k = 2 # should work for any number
dfb = dfa.copy()
dfb = dfb.sort_values(['date','size']) #actually need in dfa too
# get the k-mean
dfb['avg_price'] = dfb.groupby('date').price.rolling(k).mean().values
#to look for the k nearest sizes in merge_asof
dfb['size'] = dfb.groupby('date')['size'].rolling(k).mean().values
# add one business day to shift all the date 
dfb['date'] = dfb['date'] + pd.tseries.offsets.BDay() 
dfb = dfb.dropna().drop('price',1)
dfb['size'] = dfb['size'].astype(int) #needed for the merge_asof
print (dfb)

        date   size  avg_price
1 2018-08-02    105      232.5
2 2018-08-02    117      247.5
4 2018-08-03    115      220.0
5 2018-08-03    135      245.0

您可以按日期和最近大小使用merge_asof(sort_values对于使用方法):

dfa = (pd.merge_asof(dfa.sort_values('size'), dfb.sort_values('size'), 
                     on='size',by='date',direction='nearest')
         .sort_values(['date','size']).reset_index(drop=True))

结果是 dfa:

        date  price  size  avg_price
0 2018-08-01    220   100        NaN
1 2018-08-01    245   110        NaN
2 2018-08-01    250   125        NaN
3 2018-08-02    210   110      232.5
4 2018-08-02    230   120      247.5
5 2018-08-02    260   150      247.5
6 2018-08-03    200   115      220.0