滚动总和，周期由 Pandas Column 给出

Question

尝试在 pandas 中计算 A 列中值的滚动总和，B 列中给出回溯期，C 列中存储滚动总和的结果。

Index | Column A | Column B || Column C |
      | -------- | -------- || -------- |
  0   |     1    |    1     ||    1     |
  1   |     2    |    2     ||    3     |
  2   |     1    |    3     ||    4     |
  3   |     3    |    2     ||    4     |
  4   |     2    |    4     ||    8     |

例如，对于最后一行，滚动总和应将 A 列中的最后 4 个值相加，因为 B 列中给出了 4。

避免循环是最佳选择。

虽然是一个简单的任务，但我还没有想出解决方案。

Answer 1

由于您的滚动总和取决于所有值，因此您必须为每个窗口计算一个。这可以使用 numpy 和 索引查找来完成：

import numpy as np

idx, vals = pd.factorize(df['B'])

df['C'] = np.vstack([
    df['A'].rolling(v, min_periods=1).sum()
    for v in vals]
)[idx, np.arange(len(df))]

输出：

   A  B    C
0  1  1  1.0
1  2  2  3.0
2  1  3  4.0
3  3  2  4.0
4  2  4  8.0