pandas基于遍历行添加新列

时间:2017-05-24 14:51:23

标签: python pandas

我有一个企业的交易清单。

示例数据框:

userid  date        amt start_of_day_balance
123     2017-01-04  10  100.0
123     2017-01-05  20  NaN
123     2017-01-02  30  NaN
123     2017-01-04  40  100.0

并不总是检索天平衡开始(在这种情况下,我们会收到NaN)。但从我们知道任何一天的日均余额开始的那一刻起,我们就可以在事后的每次交易后准确估计余额。

在此示例中,新列应如下所示:

userid  date        amt start_of_day_balance calculated_balance
123     2017-01-04  10  100.0                110
123     2017-01-05  20  NaN                  170
123     2017-01-02  30  NaN                  NaN
123     2017-01-04  40  100.0                150

请注意,无法确定同一天发生的交易的确切顺序 - 在这种情况下,我很高兴忽略这一点。

我的问题是如何创建这个新专栏。类似的东西:

df['calculated_balance'] = df.sort_values(['date']).groupby(['userid'])\
['amt'].cumsum() + df['start_of_day_balance'].min()
由于NaNs,

无效。

我也不想过滤掉在第一次记录的每日余额开始之前发生的任何交易。

1 个答案:

答案 0 :(得分:0)

我想出了一个似乎有用的解决方案。我不确定它有多优雅。

def calc_estimated_balance(g):
    # find the first date which has a start of day balance
    first_date_with_bal = g.loc[g['start_of_day_balance'].first_valid_index(), 'date']

    # only calculate the balance if date is greater than or equal to the date of the first balance
    g['calculated_balance'] = g[g['date'] >= first_date_with_bal]['amt'].cumsum().add(g['start_of_day_balance'].min())
    return g


df = df.sort_values(['date']).groupby(['userid']).apply(calc_estimated_balance)