如何用单个操作同时获取Pandas分组数据的最大值及其索引？

需求说明 #

需要从按年分组的Pandas数据表中，一次性获取每年对应的预订量最大值，以及该最大值对应的月份，无需分别调用max()和idxmax()。

原始代码与问题 #

你当前的代码先统计了2016、2017年各月份的预订量：

per_2016_17_year = bookings[['arrival_date_year', 'arrival_date_month']]\
    .query('arrival_date_year in [2016,2017]')\
    .groupby(['arrival_date_year', 'arrival_date_month'])\
    .agg({ 'arrival_date_month' : 'count'})

之后只能分别获取最大值和对应月份：

# 获取2016年最大预订量
per_2016_17_year.loc[2016].max()
# 输出：arrival_date_month    6203

# 获取2016年最大预订量对应的月份
per_2016_17_year.loc[2016].idxmax()
# 输出：arrival_date_month    October

解决方案 #

可以通过以下两种方式一次性获取每年的最大值及对应月份：

方法1：通过索引定位最大行

先优化统计代码（重命名列更清晰），再按年分组找到最大行的索引，最后提取对应数据：

# 1. 统计每年各月的预订量，并重命名列
per_2016_17_year = bookings[['arrival_date_year', 'arrival_date_month']]\
    .query('arrival_date_year in [2016, 2017]')\
    .groupby(['arrival_date_year', 'arrival_date_month'])\
    .size()\
    .rename('booking_count')

# 2. 按年分组，获取每个组中预订量最大的行的索引
max_row_indices = per_2016_17_year.groupby('arrival_date_year').idxmax()

# 3. 通过索引提取对应的年份、月份和最大预订量
result = per_2016_17_year.loc[max_row_indices].reset_index()

执行后result的输出示例：

方法2：使用agg同时聚合最大值与对应月份

直接在按年分组的聚合操作中，同时计算最大值和对应的月份：

# 先完成基础统计（同方法1）
per_2016_17_year = bookings[['arrival_date_year', 'arrival_date_month']]\
    .query('arrival_date_year in [2016, 2017]')\
    .groupby(['arrival_date_year', 'arrival_date_month'])\
    .size()\
    .rename('booking_count')\
    .reset_index()

# 按年聚合，同时获取最大预订量和对应月份
result = per_2016_17_year.groupby('arrival_date_year').agg(
    max_booking=('booking_count', 'max'),
    max_month=('arrival_date_month', lambda x: x[per_2016_17_year['booking_count'] == x.max()].values[0])
)

执行后result的输出示例：

内容的提问来源于stack exchange，提问作者maximguzovguzov