Pandas 在 groupby 之后添加带有列表的计数和列

Question

I have the following dataframe df :

Month       Cat Constr  Part
1/1/2021    asd v1      p1
1/1/2021    asd v1      p2
1/1/2021    asd v2      p1
1/1/2021    pqr v1      p1
1/1/2021    pqr v1      p2
1/1/2021    pqr v2      p2

I am trying to achieve the following:

Month       Cat Constr  count   part
1/1/2021    asd v1      2       [p1, p2]
1/1/2021    asd v2      1       [p1]
1/1/2021    pqr v1      2       [p1,p2]
1/1/2021    pqr v2      1       [p2]

To get the count column I do the following:

df.groupby(['Month', 'Cat', 'Constr']).agg(['count'])

But I am not sure how to add the list of Part who are part of groupby

I am trying to do group by ['Month', 'Cat', 'Constr'] and find the number of Part for each groupby unique combination and a list of Part which form that group.

Answer 1

Use groupby as you did but apply list :

>>> df.groupby(['Month', 'Cat Constr']).agg(list).reset_index()

       Month Cat Constr      Part
0 2021-01-01     asd v1  [p1, p2]
1 2021-01-01     asd v2      [p1]
2 2021-01-01     pqr v1  [p1, p2]
3 2021-01-01     pqr v2      [p2]

To get all together:

>>> df.groupby(['Month', 'Cat Constr'])['Part'] \
      .agg(count='count', part=list)

       Month Cat Constr  count      part
0 2021-01-01     asd v1      2  [p1, p2]
1 2021-01-01     asd v2      1      [p1]
2 2021-01-01     pqr v1      2  [p1, p2]
3 2021-01-01     pqr v2      1      [p2]