如何通过某个值的总和将Pandas DataFrame拆分为多个DataFrame

Question

我有一个DataFrame，看起来像这样：

Name      Count
0   John      0.25
1   Adam      0.75
2   Michael   1.5
3   Jane      0.8
4   Anna      2.0
5   Sarah     0.25

我的目标是根据限制值将此DataFrame拆分为多个DataFrame。对于这个例子，极限值将是3，所以结果看起来像：

Name      Count
0   John      0.25
1   Adam      0.75
2   Michael   1.5

Name      Count
0  Jane      0.8
1  Anna      2.0
Name      Count
0  Sarah     0.25

重点是，在每个新的数据帧中，计数的总和接近极限值，但不超过极限值(即Jane在第二个数据帧中。因为如果她被包括在第一个数据帧，总和将是3.3，这超过了极限值3(。

我相信使用.iterrows可以做到这一点，但这相当沉重/缓慢，因此正在寻找另一种解决方案。

Answer 1

我们可以，

LIMIT = 3
assert df['Count'].le(LIMIT).all()
groups = []
sum = 0
group = 0
for val in df['Count']:
sum += val
if sum > LIMIT:
group += 1
sum = val
groups.append(group)

my_dict = {f'Group {i}' : group.reset_index(drop=True) 
for i, group in df.groupby(groups)}
#print(my_dict['Group 0'])

Answer 2

为了加快速度，我们可以使用即时编译(JIT(的numba，它将非常高效：

从我的另一个答案调整功能

from numba import njit
import numpy as np
@njit
def cumsum_reset(array, limit):
total = 0
counter = 0
groups = np.empty(array.shape[0])

for idx, i in enumerate(array):
total += i
if total >= limit:
total = 0
total += i
counter += 1
groups[idx] = counter
else:
groups[idx] = counter
return groups
groups = cumsum_reset(df['Count'].to_numpy(), 3)
for _, grp in df.groupby(groups):
print(grp)
Name  Count
0     John   0.25
1     Adam   0.75
2  Michael   1.50
Name  Count
3  Jane    0.8
4  Anna    2.0
Name  Count
5  Sarah   0.25