我可以在 Pandas 中执行动态的行累积吗？

Question

如果我有以下数据框，可以这样派生： df = pd.DataFrame(np.random.randint(0, 10, size=(10, 1)))    00   01   22   83   14   05   06   77   08   29   2有没有一种有效的方式cumsum行有限制，每次达到这个限制时，开始一个新的cumsum. 达到每个限制后（无论是多少行），都会创建一个包含总 cumsum 的行。下面我创建了一个执行此操作的函数示例，但速度非常慢，尤其是当数据帧变得非常大时。我不喜欢我的函数是循环的，我正在寻找一种使它更快的方法（我猜是一种没有循环的方法）。def foo(df, max_value):    last_value = 0    storage = []    for index, row in df.iterrows():        this_value = np.nansum([row[0], last_value])        if this_value >= max_value:            storage.append((index, this_value))            this_value = 0        last_value = this_value    return storage如果你喜欢我的函数：foo(df, 5) 在上面的上下文中，它返回：   02  106  8

Answer 1

循环不一定是坏的。诀窍是确保它在低级对象上执行。在这种情况下，您可以使用 Numba 或 Cython。例如，使用具有以下功能的生成器numba.njit：

from numba import njit

@njit

def cumsum_limit(A, limit=5):

    count = 0

    for i in range(A.shape[0]):

        count += A[i]

        if count > limit:

            yield i, count

            count = 0

idx, vals = zip(*cumsum_limit(df[0].values))

res = pd.Series(vals, index=idx)

要演示使用 Numba 进行 JIT 编译的性能优势：

import pandas as pd, numpy as np

from numba import njit

df = pd.DataFrame({0: [0, 2, 8, 1, 0, 0, 7, 0, 2, 2]})

@njit

def cumsum_limit_nb(A, limit=5):

    count = 0

    for i in range(A.shape[0]):

        count += A[i]

        if count > limit:

            yield i, count

            count = 0

def cumsum_limit(A, limit=5):

    count = 0

    for i in range(A.shape[0]):

        count += A[i]

        if count > limit:

            yield i, count

            count = 0

n = 10**4

df = pd.concat([df]*n, ignore_index=True)

%timeit list(cumsum_limit_nb(df[0].values))  # 4.19 ms ± 90.4 µs per loop

%timeit list(cumsum_limit(df[0].values))     # 58.3 ms ± 194 µs per loop

Answer 2

更简单的方法：

def dynamic_cumsum(seq,limit):

    res=[]

    cs=seq.cumsum()

    for i, e in enumerate(cs):

        if cs[i] >limit:

            res.append([i,e])

            cs[i+1:] -= e

    if res[-1][0]==i:

        return res

    res.append([i,e])

    return res

结果：

x=dynamic_cumsum(df[0].values,5)

x

>>[[2, 10], [6, 8], [9, 4]]