Группируйте столбцы фрейма данных в списках

#python #pandas #pandas-groupby

Вопрос:

У меня есть фрейм данных, который выглядит так:

 | 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 |
 

Я хочу сгруппировать его в столбце «идентификатор», чтобы каждая строка была в виде списка.

Ожидаемый Результат:

 | 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]  |
 

Объяснение:
Значения для идентификатора 1 в столбце A находятся в одном списке аналогично для других. Длина списка зависит от количества записей этого идентификатора в исходном кадре данных.

Мой подход:

 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'))
 

Хотя это дает мне желаемый результат, но он медленный для больших кадров данных, и я чувствую, что это не очень оптимально, так как мы каждый раз группируемся по идентификатору и объединяемся.
Должен быть способ, которым я один раз группируюсь по идентификатору и делаю что-то в столбцах.tolist() и он выдает тот же результат.

Любая помощь будет признательна. Спасибо

Ответ №1:

Воспользуйся GroupBy.agg :

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

Или:

 #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]
 

Комментарии:

1. Спасибо! Это отлично работает, хотя для меня было удивительно отметить, что скорость моего кода и этой агрегатной функции была в одном диапазоне и не сильно отличалась. Причиной моего удивления является тот факт, что мой код вызывает часть groupby («идентификатор») 4 раза и только один раз в этом. Есть какие-нибудь мысли о том, почему это может быть?

2. @Alex — Потому что он использует groupby.apply и обрабатывает каждый столбец в фрейме данных вместе, но groupby.agg обрабатывает каждый столбец отдельно. Это означает сначала A, затем b… и последнее D.