Эквивалент преобразования в R/ddply в Python/pandas?

В функции R ddply вы можете вычислять любые новые столбцы по группам и добавлять результат к исходному фреймворку данных, например:

ddply(mtcars, .(cyl), transform, n=length(cyl)) # n is appended to the df

В Python/ pandas я сначала вычислил его, а затем объединил, например:

df1 = mtcars.groupby("cyl").apply(lambda x: Series(x["cyl"].count(), index=["n"])).reset_index()
mtcars = pd.merge(mtcars, df1, on=["cyl"])

или что-то в этом роде.

Тем не менее, я всегда чувствую, что это довольно сложно, так что можно сделать все это раз?

Спасибо.

Ответ 1

Вы можете добавить столбец в DataFrame, присвоив ему результат операции groupby/transform:

mtcars['n'] = mtcars.groupby("cyl")['cyl'].transform('count')

import pandas as pd
import pandas.rpy.common as com

mtcars = com.load_data('mtcars')
mtcars['n'] = mtcars.groupby("cyl")['cyl'].transform('count')
print(mtcars.head())

дает

                    mpg  cyl  disp   hp  drat     wt   qsec  vs  am  gear  carb   n
Mazda RX4          21.0    6   160  110  3.90  2.620  16.46   0   1     4     4   7
Mazda RX4 Wag      21.0    6   160  110  3.90  2.875  17.02   0   1     4     4   7
Datsun 710         22.8    4   108   93  3.85  2.320  18.61   1   1     4     1  11
Hornet 4 Drive     21.4    6   258  110  3.08  3.215  19.44   1   0     3     1   7
Hornet Sportabout  18.7    8   360  175  3.15  3.440  17.02   0   0     3     2  14

Чтобы добавить несколько столбцов, вы можете использовать groupby/apply. Убедитесь, что применяемая вами функция возвращает DataFrame с тем же индексом, что и его вход. Например,

mtcars[['n','total_wt']] = mtcars.groupby("cyl").apply(
    lambda x: pd.DataFrame({'n': len(x['cyl']), 'total_wt': x['wt'].sum()},
                           index=x.index))
print(mtcars.head())

дает

                    mpg  cyl  disp   hp  drat     wt   qsec  vs  am  gear  carb   n  total_wt
Mazda RX4          21.0    6   160  110  3.90  2.620  16.46   0   1     4     4   7    21.820
Mazda RX4 Wag      21.0    6   160  110  3.90  2.875  17.02   0   1     4     4   7    21.820
Datsun 710         22.8    4   108   93  3.85  2.320  18.61   1   1     4     1  11    25.143
Hornet 4 Drive     21.4    6   258  110  3.08  3.215  19.44   1   0     3     1   7    21.820
Hornet Sportabout  18.7    8   360  175  3.15  3.440  17.02   0   0     3     2  14    55.989