Панды: как группировать (и суммировать) по категориям и сохранять информацию из подкатегории

Question

Это дополнительный вопрос к Пандам: как подгруппировать (и суммировать) верхние N наблюдений в подкатегориях? Там было продемонстрировано, как вы можете найти сумму трех верхних месяцев за каждый год. в этом фрейме данных:

Пример кадра данных

    year      month   passengers
0    1949    January         112
1    1949   February         118
2    1949      March         132
3    1949      April         129
4    1949        May         121
5    1949       June         135
.
.
.
137  1960       June         535
138  1960       July         622
139  1960     August         606
140  1960  September         508
141  1960    October         461
142  1960   November         390
143  1960   December         432

Так что в итоге вы получите:

    year  passengers
0   1949         432
1   1950         498
2   1951         582
3   1952         690
4   1953         779
5   1954         859
6   1955        1026
7   1956        1192
8   1957        1354
9   1958        1431
10  1959        1579
11  1960         176

Число 432 for 1949 является суммой 148+148+136 for the months July, August and September. Мой вопрос сейчас такой:

Можно ли выполнять те же вычисления и одновременно сохранять соответствующие подкатегории в виде списка в своем собственном столбце?

Желаемый выход

(я только проверил фактическую сумму за 1949 год. 1950 составлен):

        year  passengers  months
    0   1949         432  July, August, September 
    1   1950         498  August, September, December
    2   1951         582  .
    3   1952         690  .
    4   1953         779  .
    5   1954         859  .
    6   1955        1026  .
    7   1956        1192  .
    8   1957        1354  .
    9   1958        1431  .
    10  1959        1579  .
    11  1960         176  .

Воспроизводимый код и данные:

import pandas as pd
import seaborn as sns
df = sns.load_dataset('flights')
print(df.head())

df2 = df.groupby('year')['passengers'].apply(lambda x: x.nlargest(3).sum()).reset_index()
print(df2.head())

ДФ:

   year     month  passengers
0  1949   January         112
1  1949  February         118
2  1949     March         132
3  1949     April         129
4  1949       May         121

df2:

   year  passengers
0  1949         432
1  1950         498
2  1951         582
3  1952         690
4  1953         779

Спасибо за любые предложения!

Answer 1

В качестве альтернативы - сгруппируйте, затем примените, используя pd.DataFrame.nlargest вместо пользовательской функции / лямбды, затем перегруппируйте в индексе и примените подходящие agg с, например:

new_df = (
    df.groupby('year').apply(pd.DataFrame.nlargest, 3, 'passengers')
    .groupby(level=0).agg({'passengers': 'sum', 'month': ', '.join})
    # optionally reset index
    # .reset_index()
)

Это даст вам:

      passengers                    month
year                                     
1949         432  July, August, September
1950         498  July, August, September
1951         582  July, August, September
1952         690       August, July, June
1953         779       August, July, June
1954         859       July, August, June
...

Кажется, что year в качестве индекса имеет смысл в результирующем кадре, но примените .reset_index(), если нет.

Answer 2

Вы могли бы

In [69]: df.groupby('year').apply(lambda x: 
           x.nlargest(3, 'passengers').agg(
              {'passengers': 'sum', 'month': lambda x: ', '.join(x.values)}
             )).reset_index()
Out[69]:
    year  passengers                    month
0   1949         432  July, August, September
1   1950         498  July, August, September
2   1951         582  July, August, September
3   1952         690       August, July, June
4   1953         779       August, July, June
5   1954         859       July, August, June
6   1955        1026       July, August, June
7   1956        1192       July, August, June
8   1957        1354       August, July, June
9   1958        1431       August, July, June
10  1959        1579       August, July, June
11  1960        1763       July, August, June

Answer 3

Вот одно решение, использующее nlargest.

def largest(x, k):
    vals = x.nlargest(n=k, columns=['passengers'])
    return [vals['passengers'].sum(), vals['month'].tolist()]

g = df.groupby('year').apply(largest, k=3).reset_index()
joiner = pd.DataFrame(g[0].values.tolist(), columns=['passengers', 'months'])

res = g.drop(0, axis=1).join(joiner)

print(res)

   year  passengers               months
0  1949         382  [March, April, May]

Я намеренно сохраняю months в виде списка, вы можете преобразовать его в строку через запятую, если хотите.

Answer 4

Используйте пользовательскую функцию с GroupBy.apply, идея сначала сортируется по sort_values, а затем вызывается head для верхних значений по группам:

def f(x):
    x = x.head(3)
    names = ['passengers','months']
    return pd.Series([x['passengers'].sum(), ', '.join(x['month'])], index=names)

df2 = df.sort_values('passengers', ascending=False).groupby('year').apply(f).reset_index()
print(df2.head())
   year  passengers                   months
0  1949         432  July, August, September
1  1950         498  July, August, September
2  1951         582  July, August, September
3  1952         690       August, July, June
4  1953         779       August, July, June