PandasのDataFrameでグループ別にサンプルをN個抜き出す方法
新刊情報
技術書典8の新刊『モザイク除去から学ぶ 最先端のディープラーニング』（A4・195ページ）好評通販中です！ 機械学習の入門からGANの最先端までを書いたおすすめの本となっています！ Boothで試し読みできます。情報まとめ・質問用GitHub
「PandasでGroupbyでグルーピングしたはいんだけど、そこからグループ別にサンプルを1個、2個…と抜き出す、SQLでよくやるやつってどうやるんだっけ？」ということが気になったので、調べました。ちゃんとした方法があります。
例題
今、中国地方と四国地方の県と面積をDataFrameにしてみました。ここから、中国地方と四国地方の面積が小さい県を選んでみます。
import pandas as pd
prefs = {"pref_name":["鳥取", "島根", "岡山", "広島", "山口", "徳島", "香川", "愛媛", "高知"],
"category":["中国", "中国", "中国", "中国", "中国", "四国", "四国", "四国", "四国"],
"area":[3507, 6708, 7010, 8480, 6114, 4147, 1862, 5679, 7105]}
df = pd.DataFrame(prefs)
print(df)
pref_name category area
0 鳥取 中国 3507
1 島根 中国 6708
2 岡山 中国 7010
3 広島 中国 8480
4 山口 中国 6114
5 徳島 四国 4147
6 香川 四国 1862
7 愛媛 四国 5679
8 高知 四国 7105
グルーピングして抜き出すときはhead()
全体の流れとして、面積（area）でソートして、区分（category）でgroupbyすればいいというのはわかります。でもグルーピングした後抜き出すのはどうやるんしょう？
答えはpandas.core.groupby.GroupBy.headです。headなので、頭からサンプルをN個取ります。後ろから取りたいときはtail()を使います。
やってみましょう。
# sort
sorted = df.sort_values(["area"], 0, [True])
print(sorted)
# group by
grouped = sorted.groupby("category").head(1)
print(grouped)
pref_name category area
6 香川 四国 1862
0 鳥取 中国 3507
5 徳島 四国 4147
7 愛媛 四国 5679
4 山口 中国 6114
1 島根 中国 6708
2 岡山 中国 7010
8 高知 四国 7105
3 広島 中国 8480
pref_name category area
6 香川 四国 1862
0 鳥取 中国 3507
この通り、四国で最も狭い香川県と、中国で最も狭い鳥取県を抜き出すことができました。
後ろから抜き出したいときはtail()
さっきは面積が小さい県を抜き出しましたが、逆に最も面積が大きい県を抜き出してみましょう。sort_valuesを降順ソートにしてhead()でもいいですが、ここではtail()を使ってみます。
grouped = sorted.groupby("category").tail(1)
print(grouped)
pref_name category area
8 高知 四国 7105
3 広島 中国 8480
この通り、高知と広島が出てきました。
ソートをしてグルーピングすると順番がバラバラになるので、グループ間で直したい場合は再ソートが必要
ここがちょっと曲者ですが、head()に2以上を入れるとグループ別にバラバラの値が出てくるということがおきます。
# sort
sorted = df.sort_values(["area"], 0, [True])
# group by
grouped = sorted.groupby("category").head(2)
print(grouped)
pref_name category area
6 香川 四国 1862
0 鳥取 中国 3507
5 徳島 四国 4147
4 山口 中国 6114
なので、グループごとに一緒にして表示したいという場合はグループで再ソートする必要がありそうです（ここもっと上手いやり方あったら教えてください）。
# さらにソート
resort = grouped.sort_values(["category"], 0, [False])
print(resort)
pref_name category area
6 香川 四国 1862
5 徳島 四国 4147
0 鳥取 中国 3507
4 山口 中国 6114
このように、四国で面積の小さい2県：香川・徳島と、中国で面積の小さい2県：鳥取・山口を抽出することができました。愛媛県は山口県より面積が小さいですが（5679）、四国では3番目に小さいのでここには含まれていません。
実は公式ドキュメントに載ってた
この方法探してて、あんまいい方法ないなと思っていたら公式ドキュメントに載っていました。ただかなり下の方に載っている方法なので、気をつけてみないと見落とします。詳しく知りたかったら見てみてください。
Group By: split-apply-combine
https://pandas.pydata.org/pandas-docs/stable/groupby.html
新刊情報
技術書典8の新刊『モザイク除去から学ぶ 最先端のディープラーニング』好評通販中（A4・195ページ）です！ Boothで試し読みもできるのでよろしくね！
【一般同人】「じゅ～しぃ～すくりぷと」の『モザイク除去から学ぶ 最先端のディープラーニング』が3Fにて好評販売中！今作は誰もが夢見る"モザイク除去"を起点として機械学習・ディープラーニングの知識をたっぷり学べる注目作品。初心者向けの話題からマニアックな情報まで網羅しております！！ pic.twitter.com/HX1wAXESod
— とらのあな 秋葉原店A (@tora_akiba_a) March 8, 2020
これが君の望んでいる、「ディープラーニング－ロマン－」
— こしあん@技術書典8 Day2 あ03 委託 (@js_koshian2) March 22, 2020
誰もが夢見る「モザイク除去」から機械学習を学ぼう。ご好評につき鋭意増刷中！
Booth：https://t.co/d2oXBf0kQ6
とら：https://t.co/5P0HJ7r8sL
試し読み：https://t.co/1z2XScVNsz#技術書典 pic.twitter.com/tttO6VnTXG
- Booth（物理、電子版）：https://koshian2.booth.pm/items/1835219
- とらのあな（物理版：お急ぎの場合はとらのあなをお使いください。Boothよりも物理本の到着が早いです（即日で発送されます）。電子版プレゼントキャンペーンもあります ：https://ec.toranoana.shop/tora/ec/item/040030818462/
- メロンブックス（物理版）https://www.melonbooks.co.jp/detail/detail.php?product_id=670234（電子版とセットになります）
- 感想・情報まとめ・質問用：https://github.com/koshian2/MosaicDeeplearningBook