[pandas]groupbyの最初・最後の行を求めるfirst・last関数の話、headやnthとの違い

pandasで、ある特定の列の値に応じてグループ化(集計・集約)し、特定の列の値ごとに最初の行(もしくは最後の行)を求めたいときの話。

ある特定の列の値に応じてグループ化するにはgroupby関数を使う。 pandasのgroupby関数の返り値はGroupByオブジェクトというやつになる。(正確にはDataFrameGroupByまたはSeriesGroupByオブジェクト)GroupByオブジェクトに対して、やりたい操作に対応する関数を適用すれば結果が求まる。平均のmeanや合計のsum、最大値のmaxなどが有名だろう。

で、「グループごとの最初のレコードが欲しいんだけど、そういう関数はあるのだろうか」と思い、 "pandas groupby first"で検索したら、普通に出てきた。firstという関数があるので、それを使えばよい。

first関数の説明ページはこちら。
pandas.core.groupby.GroupBy.first — pandas 0.25.1 documentation

pandas.core.groupby.GroupBy.first
Compute first of group values.

これだけ。説明はいたって単純だ。

また、StackOverflowの質問を色々見ていたら、firstの他にnth、headを使っても同様のことができると書いてあったので、合わせて実験してみよう。
特に、NaNを含むときには挙動が複雑になり、思わぬ落とし穴にはまる可能性もあるので、詳述している。

基本的な例:グループごとに最初の行を選択する(first、nth、head)

以下、出力された結果は# ---の後に記載している。

import pandas as pd
<200b>
df = pd.DataFrame({
    'name'    : ['Alice', 'Bob', 'Charlie', 'David', 'Eve', 'Fred', 'George', 'Helen', 'Ian', 'John'],
    'class'   : ['A', 'A', 'B', 'B', 'A', 'C', 'A', 'B', 'C', 'C'],
    'English' : list(range(0, 100, 10)),
    'Math'    : list(range(100, 0, -10))
})
df

# ---

      name class  English  Math
0    Alice     A        0   100
1      Bob     A       10    90
2  Charlie     B       20    80
3    David     B       30    70
4      Eve     A       40    60
5     Fred     C       50    50
6   George     A       60    40
7    Helen     B       70    30
8      Ian     C       80    20
9     John     C       90    10
df.groupby("class")

# ---

<pandas.core.groupby.generic.DataFrameGroupBy object at 0x03DB2950>

groupbyの結果はDataFrameGroupBy オブジェクトであった。これに対して各関数を適用し、結果を見てみよう。

df.groupby("class").first()

# ---

          name  English  Math
class                        
A        Alice        0   100
B      Charlie       20    80
C         Fred       50    50

もとのデータと見比べて欲しい。各Classの最初のデータだけが抽出されたことが分かる。

df.groupby("class").nth(0)

# ---

          name  English  Math
class                        
A        Alice        0   100
B      Charlie       20    80
C         Fred       50    50

n番目の行を表示するための関数なので、n番目→n-th→nth関数という名前になったようだ。
最初の行を表示したいので、引数に0を指定する。結果はfirstの時と全く同じだ。

df.groupby("class").head(1)

# ---

      name class  English  Math
0    Alice     A        0   100
2  Charlie     B       20    80
5     Fred     C       50    50

head関数は「グループごとに最初のn行を表示する」メソッドである。今回は最初の1行だけが欲しいので引数に1を入れている。
first、nthの結果では、groupbyに用いていたclass列が行名に移動している。それに対して、headの結果ではclass列は元のままで、元のindexが行名に使われている。

基本的な例:グループごとに最後の行を選択する(last、nth、tail)

正反対の操作、つまりグループごとに最後の行を選択する操作の例である。 先ほどとほぼ同じ話なので、簡潔に述べる。

  • firstの反対がlast
  • nthは関数名は変わらない。最後の要素を指定するために引数に-1を入れる
  • headの反対がtail
df.groupby("class").last()

# ---

         name  English  Math
class                       
A      George       60    40
B       Helen       70    30
C        John       90    10
df.groupby("class").nth(-1)

# ---

         name  English  Math
class                       
A      George       60    40
B       Helen       70    30
C        John       90    10
df.groupby("class").tail(1)

# ---

     name class  English  Math
6  George     A       60    40
7   Helen     B       70    30
9    John     C       90    10

NaNを含むときの挙動

よく似た挙動をするこれらの関数だが、データにNaNが入っていた場合がなかなか厄介なようだ。
最初の行の方だけしか確認していないが……最後の行についても同様だろう。

まずNaNが混じったデータを作ります。

df = pd.DataFrame({
    'name'    : ['Alice', 'Bob', 'Charlie', 'David', 'Eve', 'Fred', 'George', 'Helen', 'Ian', 'John'],
    'class'   : ['A', 'A', 'B', 'B', 'A', 'C', 'A', 'B', 'C', 'C'],
    'English' : list(range(0, 100, 10)),
    'Math'    : list(range(100, 0, -10))    
})
import numpy as np
df.loc[0, "name"] = np.nan
df.loc[0, "English"] = np.nan
df.loc[0, "Math"] = np.nan
df.loc[1, "Math"] = np.nan
df.loc[5, "English"] = np.nan
df

# ---

      name class  English  Math
0      NaN     A      NaN   NaN
1      Bob     A     10.0   NaN
2  Charlie     B     20.0  80.0
3    David     B     30.0  70.0
4      Eve     A     40.0  60.0
5     Fred     C      NaN  50.0
6   George     A     60.0  40.0
7    Helen     B     70.0  30.0
8      Ian     C     80.0  20.0
9     John     C     90.0  10.0

単純なほうから述べる。

まずheadは、単純に先頭からN行を抜き出す動作であり、NaNであろうがなかろうが関係ない。 pandas.core.groupby.GroupBy.head — pandas 0.25.2 documentation

df.groupby("class").head(1)

# ---

      name class  English  Math
0      NaN     A      NaN   NaN
2  Charlie     B     20.0  80.0
5     Fred     C      NaN  50.0

nth

次にnthについて。
この関数にはdropnaパラメータがある。dropnaにはNone、'any'、'all'を指定できる。デフォルトはNoneだ。 pandas.core.groupby.GroupBy.nth — pandas 0.25.2 documentation

dropna=NoneはNaNかどうかを気にせずに数えて、n番目の行を抜き出す。

df.groupby("class").nth(0)

# ---

          name  English  Math
class                        
A          NaN      NaN   NaN
B      Charlie     20.0  80.0
C         Fred      NaN  50.0

dropna='any'は、どれか1つの列でもNaNがあったらその行はノーカウントとする。
下の例だとclass AとCの結果が変わっている。

df.groupby("class").nth(0, dropna="any")

# ---

          name  English  Math
class                        
A          Eve     40.0  60.0
B      Charlie     20.0  80.0
C          Ian     80.0  20.0

dropna='all'は、全ての列がNaNであったらその行はノーカウントとする。

df.groupby("class").nth(0, dropna="all")

# ---

          name  English  Math
class                        
A          NaN      NaN   NaN
B      Charlie     20.0  80.0
C         Fred      NaN  50.0

あれれ。全ての列がNaNの行が選択されている。公式ドキュメントは

If dropna, will take the nth non-null row, dropna is either ‘all’ or ‘any’; this is equivalent to calling dropna(how=dropna) before the groupby.

なので、groupnbyに使った列(class)もNaNの場合に限ってノーカウント、ということか?よく分からない。

first

最後にfirstだ。この関数はかなり奇妙な動作をしていて、落とし穴になるので要注意だ。 「列ごとに、1番目から見ていって、最初にNaNでない値を表示する」という動きになる。列ごとに全く独立であることに注意。
といってもわかりづらいので、具体例を出そう。

df.groupby("class").first()

# ---

          name  English  Math
class                        
A          Bob     10.0  60.0
B      Charlie     20.0  80.0
C         Fred     80.0  50.0

Class Cの場合、name列とMath列は1行目(5、Fred)にNaNでないデータがあるのでそれを使う。
一方、Englishの列は1行目がNaNなので、その次のデータ(8、Ian)を使う。
結果として、別々の行から継ぎはぎしたデータが返ってくるのだ!
Class Aでも同様だ。

pandas開発のGitHubの中で、関連するissueは以下。バグなのか仕様なのか分からないが、最初の報告は2014年と、ずいぶん昔からこの挙動のようだ。
BUG: groupby.first/last with nans · Issue #8427 · pandas-dev/pandas · GitHub

このStackOverflowの質問経由で、firstがNaNの時の挙動を知った。 python - pandas: how do I select first row in each GROUP BY group? - Stack Overflow


groupbyをした後に使える関数一覧は、公式ドキュメントのここに書いてある。
GroupBy — pandas 0.25.1 documentation

その中でもよく使う関数は、公式ドキュメントのGroupByのユーザーズガイドに書いてある。
Group By: split-apply-combine — pandas 0.25.2 documentation

firstに関連する質問。 python - Pandas dataframe get first row of each group - Stack Overflow

GitHubの中で、上述したfirstの挙動に関するissue一覧。 BUG: groupby.first/last with nans · Issue #8427 · pandas-dev/pandas · GitHubこれは先ほども記載したもの
groupby().first() skips NaN values · Issue #6732 · pandas-dev/pandas · GitHub

まだ調べればいろいろありそうだけど、書き続けて疲れたのでこの辺で。
それでは。