Как сгруппировать строки DataFrame в список с помощью groupby в pandas

Вы можете сделать это, используя метод groupby для группировки по интересующему вас столбцу, а затем применить apply с функцией list ко всем группам. Вот пример того, как это сделать:

In [1]: df = pd.DataFrame({'a': ['A', 'A', 'B', 'B', 'B', 'C'], 'b': [1, 2, 5, 5, 4, 6]})
        df

Out[1]: 
   a  b
0  A  1
1  A  2
2  B  5
3  B  5
4  B  4
5  C  6

In [2]: df.groupby('a')['b'].apply(list)
Out[2]: 
a
A       [1, 2]
B    [5, 5, 4]
C          [6]
Name: b, dtype: object

In [3]: df1 = df.groupby('a')['b'].apply(list).reset_index(name='new')
        df1
Out[3]: 
   a        new
0  A     [1, 2]
1  B  [5, 5, 4]
2  C        [6]

В этом примере мы сначала создаем DataFrame df, затем группируем данные по столбцу 'a' и применяем функцию list для объединения значений столбца 'b' в списки для каждой группы. В результате мы получаем новый DataFrame df1, который содержит столбец с названиями групп и объединенными списками.

Удобный способ достижения этого будет следующим:

df.groupby('a').agg({'b': lambda x: list(x)})

Обратите внимание на возможность написания пользовательских агрегатов: ссылка на источник.

Если производительность важна, можно использовать уровень NumPy:

import numpy as np

df = pd.DataFrame({'a': np.random.randint(0, 60, 600), 'b': [1, 2, 5, 5, 4, 6]*100})

def f(df):
    keys, values = df.sort_values('a').values.T
    ukeys, index = np.unique(keys, True)
    arrays = np.split(values, index[1:])
    df2 = pd.DataFrame({'a': ukeys, 'b': [list(a) for a in arrays]})
    return df2

Тесты:

In [301]: %timeit f(df)
1000 loops, best of 3: 1.64 ms per loop

In [302]: %timeit df.groupby('a')['b'].apply(list)
100 loops, best of 3: 5.26 ms per loop

Как видно из результатов тестов, пользовательская функция f на основе NumPy работает значительно быстрее, чем стандартный метод groupby в Pandas. Это связано с тем, что NumPy оптимизирован для работы с массивами данных, что позволяет обрабатывать их более эффективно. Если у вас есть большие наборы данных и производительность критична, стоит рассмотреть возможность использования подобного подхода.

Для агрегации нескольких столбцов в виде списков, вы можете использовать любой из следующих способов с методами groupby и agg:

import pandas as pd

# Настройка данных
df = pd.DataFrame({
  'a': ['A', 'A', 'B', 'B', 'B', 'C'],
  'b': [1, 2, 5, 5, 4, 6],
  'c': ['x', 'y', 'z', 'x', 'y', 'z']
})
df

Это создаст следующий DataFrame:

   a  b  c
0  A  1  x
1  A  2  y
2  B  5  z
3  B  5  x
4  B  4  y
5  C  6  z

Чтобы агрегировать несколько столбцов в виде списков, вы можете использовать один из следующих вариантов:

df.groupby('a').agg(list)
df.groupby('a').agg(pd.Series.tolist)

Это приведет к следующему результату:

           b          c
a                      
A     [1, 2]     [x, y]
B  [5, 5, 4]  [z, x, y]
C        [6]        [z]

Если вам нужно сгруппировать и получить список только для одного столбца, вы можете преобразовать groupby в объект SeriesGroupBy, а затем вызвать метод agg. Используйте:

df.groupby('a').agg({'b': list})  # 4.42 ms 
df.groupby('a')['b'].agg(list)    # 2.76 ms - быстрее

Результат будет таким:

a
A       [1, 2]
B    [5, 5, 4]
C          [6]
Name: b, dtype: object

Таким образом, вы можете эффективно агрегировать данные в нужном вам формате, используя методы groupby и agg.

Когда следует использовать agg вместо apply?

Допустим, у вас есть следующий DataFrame:

df = pd.DataFrame({'a': ['A', 'A', 'B', 'B', 'B', 'C'], 'b': [1, 2, 5, 5, 4, 6], 'c': [1, 2, 5, 5, 4, 6]})

Если вы хотите объединить несколько столбцов в список и получить результат в виде pd.DataFrame, вы можете использовать:

df.groupby('a')[['b', 'c']].agg(list)
# или
df.groupby('a').agg(list)

Если вам нужен только один столбец в виде списка с результатом в pd.Series, то используйте:

df.groupby('a')['b'].agg(list)
# или
df.groupby('a')['b'].apply(list)

Обратите внимание, что результат в форме pd.DataFrame будет примерно в 10 раз медленнее, чем результат в pd.Series, если вы агрегируете только один столбец. Поэтому в случае работы с несколькими столбцами используйте agg.