python之Pandas :如何按类别分组(和求和)并保留子类别的信息

这是 Pandas: How to subset (and sum) top N observations within subcategories? 的后续问题在那里演示了如何在此数据框中找到每年前 3 个月的总和:

示例数据框

    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 

1949 年的 432 是 7 月、8 月和 9 月的 148+148+136 的总和。 我现在的问题是:

是否可以进行相同的计算，同时将相应的子类别作为列表保留在它自己的列中？

期望的输出

(我只查了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()) 

df:

   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 

感谢您的任何建议!

请您参考如下方法:

使用自定义函数 GroupBy.apply , idea 首先按 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