python-3.x - GroupBy 但仍保留所有行

Question

我有这个数据框

data = [[1,'A','a'],
        [1,'A','b'],
        [1,'B','a'],
        [2,'A','a'],
        [2,'A','b'],
        [2,'A','c']]

df_1 = pd.DataFrame(data = data,columns = ['id','Main','sub_steps'])

输出

   id Main Sub_steps
0   1    A         a
1   1    A         b
2   1    B         a
3   2    A         a
4   2    A         b
5   2    A         c

我想 groupby (id, Main) 并仍然保留所有行
期望输出

   id Main Sub_steps      lst
0   1    A         a    [a,b]
1   1    A         b    [a,b]
2   1    B         a      [a]
3   2    A         a  [a,b,c]
4   2    A         b  [a,b,c]
5   2    A         c  [a,b,c]

如果我只是用 id 和 main 分组并压平另一行

df_1.groupby(['id','Main']).agg({'Sub_steps':list})

我会得到这个

         Sub_steps
id Main           
1  A        [a, b]
   B           [a]
2  A     [a, b, c]

Answer 1

使用 merge在列名上重命名 groupby 返回的 pd.Series与 agg :

df_1.merge(df_1.groupby(['id','Main'])['sub_steps'].agg(list).rename('lst'),
           on=['id', 'Main']))

输出:

   id Main sub_steps        lst
0   1    A         a     [a, b]
1   1    A         b     [a, b]
2   1    B         a        [a]
3   2    A         a  [a, b, c]
4   2    A         b  [a, b, c]
5   2    A         c  [a, b, c]

您可以将聚合的输出合并回原始数据帧。