Python 我可以在pandas中执行行的动态求和吗？

如果我有以下数据帧，派生如下：

df=pd.dataframe（np.random.randint（0,10，size=（10,1））

是否有一种有效的方法

cumsum

有限制的行，并且每次达到该限制时，启动一个新的

cumsum

。达到每个限制后（无论行数多少），将创建一行，其中包含总计总和

下面我创建了一个函数的示例，它可以实现这一点，但速度非常慢，尤其是当数据帧变得非常大时。 我不喜欢我的函数是循环的，我正在寻找一种使它更快的方法（我想是一种没有循环的方法）

如果您像这样使用我的函数：

foo（df，5）

在上述上下文中，它返回：

   0
2  10
6  8

---

循环不一定是坏的。诀窍是确保它是在低级对象上执行的。在这种情况下，可以使用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)

from numba import njit, prange

@njit
def dynamic_cumsum(seq, index, max_value):
    cumsum = []
    running = 0
    for i in prange(len(seq)):
        if running > max_value:
            cumsum.append([index[i], running])
            running = 0
        running += seq[i] 
    cumsum.append([index[-1], running])

    return cumsum

要演示使用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

---

循环无法避免，但可以使用

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)

from numba import njit, prange

@njit
def dynamic_cumsum(seq, index, max_value):
    cumsum = []
    running = 0
    for i in prange(len(seq)):
        if running > max_value:
            cumsum.append([index[i], running])
            running = 0
        running += seq[i] 
    cumsum.append([index[-1], running])

    return cumsum

这里需要索引，假设您的索引不是数字/单调递增的

%timeit foo(df, 5)
1.24 ms ± 41.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit dynamic_cumsum(df.iloc(axis=1)[0].values, df.index.values, 5)
77.2 µs ± 4.01 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

---

如果索引为

Int64Index

类型，可以将其缩短为：

@njit
def dynamic_cumsum2(seq, max_value):
    cumsum = []
    running = 0
    for i in prange(len(seq)):
        if running > max_value:
            cumsum.append([i, running])
            running = 0
        running += seq[i] 
    cumsum.append([i, running])

    return cumsum

lst = dynamic_cumsum2(df.iloc(axis=1)[0].values, 5)
pd.DataFrame(lst, columns=['A', 'B']).set_index('A')

    B
A    
3  10
7   8
9   4

---

njit

功能性能

perfplot.show(
    setup=lambda n: pd.DataFrame(np.random.randint(0, 10, size=(n, 1))),
    kernels=[
        lambda df: list(cumsum_limit_nb(df.iloc[:, 0].values, 5)),
        lambda df: dynamic_cumsum2(df.iloc[:, 0].values, 5)
    ],
    labels=['cumsum_limit_nb', 'dynamic_cumsum2'],
    n_range=[2**k for k in range(0, 17)],
    xlabel='N',
    logx=True,
    logy=True,
    equality_check=None # TODO - update when @jpp adds in the final `yield`
)

日志图显示，对于较大的输入，发电机功能更快：

一种可能的解释是，随着N的增加，在

dynamic_cumsum2

中添加到不断增长的列表的开销变得显著。而

cumsum\u limit\u nb

只需

产生

更简单的方法：

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]]

---

预期结果是10、8、4吗？还有，你对指数有多挑剔？是的，10，8，4比我的10，8会更好。索引应基于最后一个值。对于10、8、4的情况，应分别为2、6、9。我在这里遇到的瓶颈是速度：/出于兴趣，对于较大的数据帧，您看到了什么样的时间安排。。对于我答案中的数据，我认为生成器的速度快了3倍。事实上，我对此感到惊讶（很高兴），因为我认为生成器的开销很大。@jpp既然你问了，请给我一点时间，我会看看是否可以使用perfplot生成timeits。就性能而言，我不会以这样或那样的方式争论，两种解决方案都非常好，坦率地说，您的发电机解决方案在我看来很简单。没问题，如果太麻烦，不需要perfplot，我只是对numba如何/如果能优化发电机感兴趣！超级有用的帖子（你们两个）！这种类型的计算在这里经常出现，因此最好知道如何有效地处理它：D@Newskooler始终使用数组或列表。否则对numba不起作用。您可以调用Series对象上的.values来访问基础numpy数组。

x=dynamic_cumsum(df[0].values,5)
x
>>[[2, 10], [6, 8], [9, 4]]