python - 如何根据特定的列组向 DataFrame 添加额外的总和列？

Question

在这种情况下，我有类似的 DataFrame

col1  col2
   a     1
   a     2
   a     3
   b     1
   b     2

我想要的是先按 col1 分组，然后对组的 col2 列求和，最后将 sum 添加到 DataFrame 并得到

col1  col2  sum
   a     1    6
   a     2    6
   a     3    6
   b     1    3
   b     2    3

Answer 1

选项 1
transform返回与原始对象具有相同索引的结果。
我使用 assign 返回带有新列的数据框副本。
参见 split-apply-combine文档以获取更多信息。

df.assign(Sum=df.groupby('col1').col2.transform('sum'))

  col1  col2  Sum
0    a     1    6
1    a     2    6
2    a     3    6
3    b     1    3
4    b     2    3

---

选项 2
使用 join基于正常 groupby 和 sum 的结果。

df.join(df.groupby('col1').col2.sum().rename('Sum'), on='col1')

  col1  col2  Sum
0    a     1    6
1    a     2    6
2    a     3    6
3    b     1    3
4    b     2    3

---

选项 3
pd.factorize 的创意方法和 np.bincount

f, u = df.col1.factorize()
df.assign(Sum=np.bincount(f, df.col2).astype(df.col2.dtype)[f])

  col1  col2  Sum
0    a     1    6
1    a     2    6
2    a     3    6
3    b     1    3
4    b     2    3