python - 按年份和 Pandas 中的 id 求和

Question

我想了解使用 pandas 在 Python 2.7 中复制以下 Stata 命令的最紧凑方法:egen yr_id_sum = total(var_to_sum)，缺失(id year)。

我想在此表中生成 yr_id_sum 列:

id    year    value  yr_id_sum
1     2010    1      3
1     2010    2      3
1     2011    3      7
1     2011    4      7
2     2010    11     23
2     2010    12     23
2     2011    13     27
2     2011    14     27

我可以按如下方式为一个分组变量执行此操作(这可能有助于阐明我正在尝试执行的操作):

def add_mean(grp):
    grp['ann_sum'] = grp['var_to_sum'].sum()
    return grp

df=df.groupby('year').apply(add_sum)

这等同于 egen year_sum = total(var_to_sum), missing by(year)。

我很难扩展像 this 这样的答案关于在我的案例中使用带有多索引的总和。

df.set_index(['year', 'id'], inplace=True)
df=df.groupby(['year', 'id').apply(add_sum)

似乎它应该做我想做的事...但我得到了 Exception: cannot handle a non-unique multi-index!

以下是我已经看过的一些答案:

• 关于 applying a user defined function to each subgroup of a Group By in Pandas 的问题接近我要找的东西。
• 我正在尝试关注 this question , 无条件支付。

Answer 1

要重现您想要的输出，您可以使用 transform:它获取 groupby 操作的结果并将其广播回原始索引。例如:

>>> df["yr_id_sum"] = df.groupby(["id", "year"])["value"].transform(sum)
>>> df
   id  year  value  yr_id_sum
0   1  2010      1          3
1   1  2010      2          3
2   1  2011      3          7
3   1  2011      4          7
4   2  2010     11         23
5   2  2010     12         23
6   2  2011     13         27
7   2  2011     14         27

基本上是

>>> df.groupby(["id", "year"])["value"].sum()
id  year
1   2010     3
    2011     7
2   2010    23
    2011    27
Name: value, dtype: int64

但重复以匹配用作索引的原始列。