是否有一种方法可以计算满足条件/阈值后重置的 Pandas 中的累计总和？

Question

我有一个特定的用例，我需要计算一个运行总计，每当它超过同一数据帧中的另一个运行总计时，它就会被重置/调整。运行总计需要减少超过其他列运行总计的数量。

这在 Excel 中可以相对容易地实现，但在 Pandas 或 Numpy 中如何执行它真的很困难。

在下面的示例中，您可以看到我想要的结果。

您可以看到“E”和“F”列是如何在下面的公式中得出的。

这是 Pandas 中的数据框。我确实尝试创建一个 lambda 函数以及使用 pandas.DataFrame.shift 方法但没有成功。

import pandas as pd

data = {'date': ['2023-01-01', '2023-01-02', '2023-01-03', '2023-01-04','2023-01-05','2023-01-06'],
        'a': [10, 10, 15, 15,15,15],
       'b': [0,5,60,5,0,80]
       }
df = pd.DataFrame.from_dict(data)

df['a rt'] = df['a'].cumsum()
df['b rt'] = df['b'].cumsum()

Answer 1

让我们定义一个自定义函数来复制 excel 公式的行为。您也可以选择使用

numba

来编译函数，以达到

C

之类的速度

from numba import njit

@njit
def func(A, B):
    a_cs, b_cs, delta = 0, 0, 0
    for a, b in zip(A, B):
        a_cs += a
        b_cs += b

        cond = b_cs > a_cs
        if cond:
            delta = b_cs - a_cs

        yield (a_cs, b_cs, delta)
        if cond:
            b_cs = a_cs # reset
        
        delta = 0


df[['a_rt', 'b_rt', 'delta']] = [*func(df.a.values, df.b.values)]

结果

         date   a   b  a_rt  b_rt  delta
0  2023-01-01  10   0    10     0      0
1  2023-01-02  10   5    20     5      0
2  2023-01-03  15  60    35    65     30
3  2023-01-04  15   0    50    35      0
4  2023-01-05  15   0    65    35      0
5  2023-01-06  15  80    80   115     35

Answer 2

这不容易向量化，因为您的值取决于先前的计算值。

然而，您可以制作自定义函数/生成器来计算这些列值：

def get_D_E_F(B, C):
    B, C = iter(B), iter(C)
    prev_d, prev_e = next(B), next(C)
    prev_f = prev_e - prev_d if prev_e > prev_d else 0
    yield prev_d, prev_e, prev_f

    while True:
        try:
            current_d = next(B) + prev_d
            current_e = next(C) + prev_e - prev_f
            current_f = current_e - current_d if current_e > current_d else 0
            yield current_d, current_e, current_f
            prev_d, prev_e, prev_f = current_d, current_e, current_f
        except StopIteration:
            break

df[['a rt', 'b rt adj.', 'delta']] = list(get_D_E_F(df['a'], df['b']))
print(df)

印花：

         date   a   b  a rt  b rt adj.  delta
0  2023-01-01  10   0    10          0      0
1  2023-01-02  10   5    20          5      0
2  2023-01-03  15  60    35         65     30
3  2023-01-04  15   0    50         35      0
4  2023-01-05  15   0    65         35      0
5  2023-01-06  15  80    80        115     35