在数据框中向量化折叠过程

Question

假设我们有一个如下所示的示例数据框：

df = pd.DataFrame({'A': [np.nan, 0.5, 0.5, 0.5, 0.5], 
                   'B': [np.nan, 3, 4, 1, 2], 
                   'C': [10, np.nan, np.nan, np.nan, np.nan]})
>>> df
     A    B     C
0  NaN  NaN  10.0
1  0.5  3.0   NaN
2  0.5  4.0   NaN
3  0.5  1.0   NaN
4  0.5  2.0   NaN

Col 'D' 通过以下操作计算：

>>> df
     A    B     C     D
0  NaN  NaN  10.0  10.0 
1  0.5  3.0   NaN   8.0 = (10 x 0.5) + 3
2  0.5  4.0   NaN   8.0 = (8 x 0.5) + 4 
3  0.5  1.0   NaN   5.0 = (8 x 0.5) + 1 
4  0.5  2.0   NaN   4.5 = (5 x 0.5) + 2 

计算 col 'D' 反映了一个折叠过程，在每次行操作中调用 col 'C' 的前一行以及 col 'A' 和 'B' 的当前行。

我尝试过使用 for 循环、functools.reduce() 和迭代器来执行此操作，但我想知道是否有另一种方法尽可能使用矢量化，以便使此操作在更大的数据集中更有效。

Answer 1

我不知道纯矢量化的 pandas/numpy 解决方案，但你可以尝试使用 numba 来加速计算：

from numba import njit


@njit
def calculate(A, B, starting_value=10):
    out = np.empty_like(A, dtype=np.float64)

    out[0] = starting_value

    for i, (a, b) in enumerate(zip(A[1:], B[1:]), 1):
        out[i] = (out[i - 1] * a) + b

    return out


df["D"] = calculate(df["A"].values, df["B"].values, 10)
print(df)

打印：

     A    B     C     D
0  NaN  NaN  10.0  10.0
1  0.5  3.0   NaN   8.0
2  0.5  4.0   NaN   8.0
3  0.5  1.0   NaN   5.0
4  0.5  2.0   NaN   4.5