累积总和但有条件地排除较早的行

Question

我有一个像这样的数据帧：

df = pd.DataFrame({
  'val_a': [3, 3, 3, 2, 2, 2, 1, 1, 1],
  'val_b': [3, np.nan, 2, 2, 2, 0, 1, np.nan, 0],
  'quantity': [1, 4, 2, 8, 5, 7, 1, 4, 2]
})

它看起来像这样：

|    |   val_a |   val_b |   quantity |
|---:|--------:|--------:|-----------:|
|  0 |       3 |       3 |          1 |
|  1 |       3 |     nan |          4 |
|  2 |       3 |       2 |          2 |
|  3 |       2 |       2 |          8 |
|  4 |       2 |       2 |          5 |
|  5 |       2 |       0 |          7 |
|  6 |       1 |       1 |          1 |
|  7 |       1 |     nan |          4 |
|  8 |       1 |       0 |          2 |

它是由 订购的val_a。我想quantity对每个val_a. 所以：

|    |   val_a |   val_b |   quantity |
|---:|--------:|--------:|-----------:|
|  0 |       3 |       3 |          1 |
|  1 |       3 |     nan |          4 |
|  2 |       3 |       2 |          2 |
|  3 |       2 |       2 |          8 |
|  4 |       2 |       2 |          5 |
|  5 |       2 |       0 |          7 |
|  6 |       1 |       1 |          1 |
|  7 |       1 |     nan |          4 |
|  8 |       1 |       0 |          2 |

这使

|   val_a |   quantity |
|--------:|-----------:|
|       3 |          7 |
|       2 |         27 |
|       1 |         34 |

然而，这是棘手的部分。

我想排除值val_b大于 key 的行val_a。我会用一个例子来澄清：

• 计算总用于当当val_a是3，没有行的具有val_b大于val_a。所以当val_a是的累计总数3是7;
• 计算总用于当时val_a是2，则行0具有val_b大于2。该行有quantity1。所以，不包括该行，当val_a是的累计总数2是27 - 1，即26；
• 计算总用于当时val_a是1，则行0，2，3，4具有val_b大于1，。该行有quantity1。所以，不包括该行，当val_a是的累计总数1是34 - 1 - 2 - 8 - 5，即18；

这是所需的输出：

|   val_a |   quantity |
|--------:|-----------:|
|       3 |          7 |
|       2 |         26 |
|       1 |         18 |

Answer 1

在 NumPy 的帮助下：

# sum without conditions
raw_sum = df.groupby("val_a", sort=False).quantity.sum().cumsum()

# comparing each `val_b` against each unique `val_a` via `gt.outer`
sub_mask = np.greater.outer(df.val_b.to_numpy(), df.val_a.unique())

# selecting values to subtract from `quantity` and summing per `val_a`
to_sub = (sub_mask * df.quantity.to_numpy()[:, np.newaxis]).sum(axis=0)

# subtracting from the raw sum
result = raw_sum - to_sub

要得到

>>> result.reset_index()

   val_a  quantity
0      3         7
1      2        26
2      1        18