根據累積總和將數據分成批次

Question

我正在嘗試根據start_date和end_date批量處理一些數據，這些數據的條件是累積總和 <= 500000。假設我有一個包含兩列的簡單數據框：

index  Date            num_books
0      2021-01-01       200000
1      2021-01-02       240000
2      2021-01-03       55000
3      2021-01-04       400000
4      2021-01-05       80000
5      2021-01-06       100000

我需要對num_books中的值進行累積總和，直到它 <= 500000 並記錄開始日期、結束日期和累積值。 這是我想要實現的一個例子

start_date  end_date   cumsum_books
2021-01-01  2021-01-03  495000
2021-01-04  2021-01-05  480000
2021-01-06  2021-01-06  100000

有沒有一種有效的方法/功能來實現這一目標？ 謝謝！

Answer 1

這是一種方法：

from io import StringIO as sio
d = sio("""
index  Date            num_books
0      2021-01-01       200000
1      2021-01-02       240000
2      2021-01-03       55000
3      2021-01-04       400000
4      2021-01-05       80000
5      2021-01-06       100000
""")

import pandas as pd
df = pd.read_csv(d, sep='\s+')

batch_num = 5*10**5
df['batch_num'] = df['num_books'].cumsum()//batch_num 

result = df.groupby('batch_num').agg(start_date=('Date', 'min'), end_date=('Date', 'max'), cumsum_books=('num_books','sum'))
print(result)

#           start_date    end_date  cumsum_books
#batch_num                                      
#0          2021-01-01  2021-01-03        495000
#1          2021-01-04  2021-01-05        480000
#2          2021-01-06  2021-01-06        100000

請注意， result dataframe 還包含超過500_000的條目，但刪除/過濾掉它是微不足道的。