在熊猫中按组顺序计算差异

Question

我正在尝试按顺序尽可能高效地计算两列之间的差异。我的 DataFrame 如下所示：

category   sales   initial_stock
    1        2         20
    1        6         20
    1        1         20
    2        4         30
    2        6         30
    2        5         30
    2        7         30

我想计算一个变量final_stock，像这样：

category   sales   initial_stock    final_stock
    1        2         20               18
    1        6         20               12
    1        1         20               11
    2        4         30               26
    2        6         30               20
    2        5         30               15
    2        7         30               8

因此，对于每个类别，final_stock 首先等于 initial_stock - sales，然后它等于 final_stock.shift() - sales。我设法用for 循环来做到这一点，但它很慢，我的感觉是可能有一个或两个线性解决方案来解决这个问题。你有什么想法吗？

谢谢

Answer 1

在“sales”上使用groupby 和cumsum 来获得每个类别的累计销售库存，然后从“initial_stock”中减去：

df['final_stock'] = df['initial_stock'] - df.groupby('category')['sales'].cumsum()
df
   category  sales  initial_stock  final_stock
0         1      2             20           18
1         1      6             20           12
2         1      1             20           11
3         2      4             30           26
4         2      6             30           20
5         2      5             30           15
6         2      7             30            8