在有条件的情况下重新启动 Pandas 中的 cumsum

问题描述

我有列金额和资产。列目标应该是金额的总和,但如果之前资产等于零,则总和应重置为当前金额。

示例:

+--------+--------+--------+
| amount | assets | target |
+--------+--------+--------+
|      6 |     10 |      6 |
|      8 |     20 |     14 |
|     -1 |      0 |     13 |
|      6 |      1 |      6 |
|     -7 |      0 |     -1 |
|      2 |      4 |      2 |
|     -5 |      7 |     -3 |
|      3 |      9 |      0 |
|      7 |      0 |      7 |
|      9 |      2 |      9 |
|      1 |      3 |     10 |
|     -4 |      5 |      6 |
+--------+--------+--------+

解决方法

GroupBy.cumsum 与通过比较列创建的组使用 0 并移位 Series.shift,处理第一个 NaNSeries.cumsum

g = df['assets'].eq(0).shift().bfill().cumsum()
#alternative
#g = df['assets'].eq(0).shift(fill_value=0).cumsum()
df['new'] = df.groupby(g)['amount'].cumsum()

print (df)
    amount  assets  target  new
0        6      10       6    6
1        8      20      14   14
2       -1       0      13   13
3        6       1       6    6
4       -7       0      -1   -1
5        2       4       2    2
6       -5       7      -3   -3
7        3       9       0    0
8        7       0       7    7
9        9       2       9    9
10       1       3      10   10
11      -4       5       6    6