python - 根据一列的总和将 Pandas 数据帧分成更小的 block

Question

所以我有一个大型 Pandas 数据框，其结构如下:

  x     y   count
blah  blah    4
blah  blah    12
blah  blah    15
...

计数列的最终总和约为 48,000。我想将整个数据帧分成大约十二个不同的 block 。此拆分的条件是我希望该 block 中的列数约为 4,000。因此每个 block 的长度可以不同，但​​计数列的总和必须在 4,000 左右。

有没有一种快速有效的方法来实现这一点？我当前的尝试涉及迭代数据帧并追加，直到行数总和超过 4,000，但速度相当慢:

df_sum = 4000

new_df = pd.DataFrame(columns=['x','y','count'])

i = 1

for index, row in df.iterrows():

     appended_sum = 0

     if appended_sum < df_sum:

        appended_sum += row['count']

        a_df.loc[len(a_df)] = row

    else:

        a_df.to_csv('split_'+str(i)+'.csv', index=False)

        i+=1

        a_df= a_df.iloc[0:0]
        appended_sum = 0

Answer 1

从这个有限的例子开始:

      x     y  count
0  blah  blah      4
1  blah  blah     12
2  blah  blah     15
0  blah  blah      4
1  blah  blah     12
2  blah  blah     15
0  blah  blah      4
1  blah  blah     12
2  blah  blah     15

这可行:

df_sum = 30
df['total'] = df['count'].cumsum()
df['batch'] = (df['total'] / df_sum).astype(int)
grouped = df.groupby('batch')
for group in grouped.groups.keys():
    data = grouped.get_group(group)
    # do something with the data