Question

给定一个熊猫数据框，当上一行的列值之和（或当前行的列值）之和等于或大于n时，我想获取每一行的索引，然后总和重新回到零。因此，例如，如果我们的数据框具有值：

index   colB
1       10
2       20
3       5
4       5
5       15
6       5
7       7
8       3

并说n = 10，那么我想要的索引为[1、2、4、5、7]，因为ColB的前一行（或当前行）总计为10。

到目前为止，我可以在此数据帧上执行for循环以获得所需的索引，但是当有很多行时，它非常慢。因此，我正在寻求有关更快方法的帮助。谢谢！

Answer 1

使用cumsum()的某种组合可能有一种巧妙的方法，但这是一个棘手的问题，因为在总和大于n之后需要重置该值。因此，这有点像滚动总和，其窗口不大于n。

我可能会为此使用自定义函数。

import pandas as pd

def go(s, n=10):
    r = []
    c = 0
    s = s.tolist()
    cv = 0
    for v in s:
        if v >= n:
            r.append(c)
            c += 1
        else:
            cv += v
            if cv >= n:
                r.append(c)
                c += 1
                cv = 0
            else:
                r.append(c)

    return r

df = pd.DataFrame.from_dict({
    'colB': {0: 10, 1: 20, 2: 5, 3: 5, 4: 15, 5: 5, 6: 1, 7: 3, 8: 1, 9: 12}})

df['group'] = df.apply(go, args=(10,))

indices = df['group'].drop_duplicates().index

请注意，我稍微修改了示例编号。

如果df是这样的：

indices是：

Int64Index([0, 1, 2, 4, 5, 9], dtype='int64')

每次满足条件时进行熊猫采样

1 个答案: