向前填充和回填 groupby 的更快方法

Question

我想在 groupby 之后ffill添加一个特定的列。bfill

我的解决方案有效：

import numpy as np
import pandas as pd

df = pd.DataFrame({
    "A": [1, 1, 1, 1, 2, 2, 2, 2],
    "B": [np.nan, 'f1', 'b1', np.nan, np.nan, 'f2', 'b2', np.nan]
})
df['B'] = df.groupby('A')['B'].apply(lambda _: _.ffill().bfill())

所以这：

    A   B
0   1   NaN
1   1   f1
2   1   b1
3   1   NaN
4   2   NaN
5   2   f2
6   2   b2
7   2   NaN

就变成这样了：

    A   B
0   1   f1
1   1   f1
2   1   b1
3   1   b1
4   2   f2
5   2   f2
6   2   b2
7   2   b2

请注意，我想要填充和 bfill 的序列将始终采用这种格式 ( Nan, x, y, Nan)

虽然这有效，但在大型数据帧上速度非常慢。

我正在寻找一些优化来使其更快（理想情况下不诉诸使用 Dask 或多处理），也许我可以进行 Numpy 优化？

我没有太多运气查看其他答案，比如这个。

Answer 1

如果你想要速度，避免 groupby 并使用 numpy 而不是 pandas 是需要遵循的好规则。这通常是不可能的，但在这里你有非常规则的数据的特殊情况，你所需要的只是形式的下标三元组[start:end:stride]：

df.iloc[0::4,1] = df.iloc[1::4,1].values
df.iloc[3::4,1] = df.iloc[2::4,1].values

说明：大多数人都知道您可以使用形式的下标[start:stop]，但您也可以添加可选stride参数。所以第一行说用元素 1,5,9,... 替换元素 0,4,8,... “值”对于删除 pandas 索引是必要的，这实际上是有害的。

通过避免 groupby 应该会更快一些。为了提高速度，您可以将 B 列输出到 numpy，在 numpy 中进行工作（基本上相同的代码），然后重新导入到 pandas：

arr = df.B.values
arr[0::4] = arr[1::4]  
arr[3::4] = arr[2::4]
df.B = arr

如果你想留在 pandas 中，你可以做的另一件事是取消堆叠，复制整个列，然后重新堆叠。无论如何，这本质上就是上面的代码所做的事情。老实说，对于这样一个矩形问题，任何数组式的方法都会相当快。