Pandas:根据日期范围和ID在Dataframe中添加金额

时间:2017-02-25 03:17:34

标签: python pandas

我问了一段时间(here)关于如何通过比较另一个数据帧中的日期来查找数据框中列的总和的问题。但是,我现在想要做类似的事情但包含一个ID可能与另一个ID具有完全匹配的日期。

这是我的代码:

df_a = pd.DataFrame({
        'end':pd.to_datetime(['1/15/2016','1/15/2016','3/15/2016','5/15/2016','5/15/2016','7/15/2016']),
        'ID':[1,2,1,2,1,1]
    })

df_a['start'] = df_a.groupby('ID')['end'].shift().fillna(0)
df_a = df_a[['start','end','ID']]

df_b = pd.DataFrame({
        'date':pd.to_datetime(['1/1/2016','1/1/2016','2/1/2016','2/1/2016','3/1/2016','3/1/2016','6/1/2016','6/1/2016','7/1/2016','7/1/2016','8/1/2016']),
        'ID':       [1, 2, 1,  2,  1,   2,   2,    1,    1,     2,     2],
        'amount':   [1, 2, 10, 20, 100, 200, 2000, 1000, 10000, 20000, 200000]
    })

我想要的输出:

         start        end   ID   amount
0   1970-01-01 2016-01-15    1        1
1   1970-01-01 2016-01-15    2        2
2   2016-01-15 2016-03-15    1      110
3   2016-01-15 2016-05-15    2      220
4   2016-03-15 2016-05-15    1        0
5   2016-05-15 2016-07-15    1    11000

我已经尝试merge()merge_asof()combine_first()groupby()并且已经接近,但并非一直如此。

这是一个非熊猫版本,但是对于大型数据集,我想这会很慢:

amount = []
for s, e, i in zip(df_a['start'], df_a['end'], df_a['ID']):
    amount.append(df_b['amount'][(s < df_b['date']) & (df_b['date'] <= e) & (df_b['ID'] == i)].sum())

df_a['amount'] = pd.Series(amount)

希望能提前得到一些帮助。

2 个答案:

答案 0 :(得分:2)

好吧,我想我根据我在原始问题中分享的链接回答了我自己的问题(感谢@piRSquared),但是必须添加一些代码。我猜这不是最有效的方法,并且有兴趣获得其他想法。

# Merge DataFrames, find date ranges, and add amounts    
df_c = pd.merge_asof(df_b, df_a, left_on = 'date', right_on = 'start', by = 'ID') \ 
         .query('date <= end').groupby(['end','ID'])['amount'].sum().reset_index()

# But that leaves out ranges for which there is no data
# so need to merge back in the original dates and fill NaNs with 0    
df_c = df_a.merge(df_c, how = 'outer').fillna(0)

       start        end   ID   amount
0 1970-01-01 2016-01-15    1      1.0
1 1970-01-01 2016-01-15    2      2.0
2 2016-01-15 2016-03-15    1    110.0
3 2016-01-15 2016-05-15    2    220.0
4 2016-03-15 2016-05-15    1      0.0
5 2016-05-15 2016-07-15    1  11000.0

答案 1 :(得分:2)

你的答案非常好。我喜欢这样,因为reindex使用fill_value会保留int dtype

mux = pd.MultiIndex.from_arrays(df_a.values.T, names=df_a.columns)

kws = dict(
    left_on='date', right_on='start',
    allow_exact_matches=True, by='ID')
mrg = pd.merge_asof(df_b, df_a, **kws).query('date <= end')
grp = mrg.groupby(['start', 'end', 'ID']).amount.sum()
grp.reindex(mux, fill_value=0).reset_index()

       start        end  ID  amount
0 1970-01-01 2016-01-15   1       1
1 1970-01-01 2016-01-15   2       2
2 2016-01-15 2016-03-15   1     110
3 2016-01-15 2016-05-15   2     220
4 2016-03-15 2016-05-15   1       0
5 2016-05-15 2016-07-15   1   11000