对列表中的数据框列进行分组

Question

I have a dataframe which looks like this:

| id | A  | B  | C  | D  |
| 1  | 50 | 51 | 52 | 53 |
| 2  | 70 | 71 | 72 | 73 |
| 1  | 80 | 81 | 82 | 83 |
| 1  | 90 | 91 | 92 | 93 |
| 2  | 40 | 41 | 42 | 43 |

I want to group it on 'id' column so that each row is in form of list.

Expected Output:

| id |      A     |      B     |      C     |      D     |
| 1  | [50,80,90] | [51,81,91] | [52,82,92] | [53,83,93] |
| 2  |   [70,40]  |   [71,41]  |   [72,42]  |   [73,43]  |

Explaination: The values for id 1 in column A are all in one list similarly for other. The length of list depends on number of records of that id in initial dataframe.

My approach:

df_grouped = df.groupby(['id'])['A'].apply(lambda x: list(x)).reset_index().merge(df.groupby(['id'])['B'].apply(lambda x: list(x)).reset_index().merge(df.groupby(['id'])['C'].apply(lambda x: list(x)).reset_index().merge(df.groupby(['id'])['D'].apply(lambda x: list(x)).reset_index()),on=['id'],how='left'))

Although this gives me the desired output, but its slow for large dataframes and I feel this is not very optimal as well as we are grouping on id each time and merging. There should be a way wherein I group on id once and do something columns.tolist() and it gives the same output.

Any help would be appreciated. Thanks

Answer 1

Use GroupBy.agg :

#all columns without id
df_grouped = df.groupby('id').agg(list).reset_index()

Or:

#columns specified in list
df_grouped = df.groupby('id')[['A','B','C','D']].agg(list).reset_index()

print (df_grouped)
   id             A             B             C             D
0   1  [50, 80, 90]  [51, 81, 91]  [52, 82, 92]  [53, 83, 93]
1   2      [70, 40]      [71, 41]      [72, 42]      [73, 43]