pandasのgroupbyを使ってグループ内で最大値（最小値）を持つ行を取得する

はじめに

データベースのカラムで、カテゴリ変数を値に持つカラムがあるとします。そのカテゴリごとの最大値もしくは最小値の列を求めたいことがあります。
pandasのidxmax()もしくはidxmin()で実現できます。

参考

• pandasで最大値・最小値の行名・列名を取得するidxmax, idxmin
• Pandas：グループ毎に括って最大の値を含む列を抜き出す

環境

• Google Colaboratory
• pandas==0.22.0

手順

タイタニック号の生死者データを使用します。
最大値の例を以下に示します。最小値の場合はidxmax()をidxmin()にすればよいです。
NaNはないとして、また、最大値（最小値）の重複はないものとします。

import pandas as pd
import seaborn as sns

df = sns.load_dataset("titanic") # タイタニック号の生死者データ
df = df.dropna() # 今回は本筋に関係ないのでNaNを落とします
df.reset_index(drop=True,inplace=True)

'embarked'は三つのカテゴリ（C,Q,S）を持ちます。これでグループ化して、それぞれのグループの'age'が最大の行を抜き出します。

ddf = df.groupby('embarked')
df.loc[ddf['age'].idxmax(),:]

df.groupby('embarked')でグループ化します。グループ化したデータフレームの'age'列からidxmax()で、それぞれのグループの最大値のインデックスを取得します。そのインデックスの行をdf.locで取得します。
idxmax()の挙動としては上から検索して、早く見つかった最大値を採用してるっぽいです。

別解

Pandas：グループ毎に括って最大の値を含む列を抜き出すに紹介されている方法だとより柔軟に書けます。
以下は自分向けに少し改変しました。

def select(df,**kwargs): # colに列名,kindに最大もしくは最小

    if kwargs['kind'] == 'min':
        val_r = df[df[kwargs['col']] == min(df[kwargs['col']]) ]
    elif kwargs['kind'] == 'max':
        val_r = df[df[kwargs['col']] == max(df[kwargs['col']]) ]
    else:
        raise Exception("パラメータ不正")

    # 全く同じ行があった場合は削除
    val_r = val_r.drop_duplicates()

    return val_r

# 最大値の行を抜き出す
df2 = (df.groupby(['embarked'])
       .apply(select,col='age',kind='max')
      )
df2 = df2.reset_index(level=1).reset_index(drop=True)
df2.index = df2['level_1'].values
df2.drop('level_1', axis=1, inplace=True)

applyするとインデックスがカテゴリになるので、それを元のインデックスにするために処理が必要になります。
また、applyは少々遅いので、大規模データだと時間がかかってしまうのが難点です。

おわりに

pandasの操作をなるべく高速に実行したい。