Python: pandas の DataFrameGroupBy#agg() には関数も渡せる

今回は pandas で DataFrame#groupby() したときに得られるオブジェクト DataFrameGroupBy が持つメソッド agg() について。これまであんまり使ってこなかったけど、関数が渡せることを知って色々と便利に使えそうだなと感じた。ちょっと前置きが長くなるので知っているところに関しては飛ばしながら読むと良いかも。

使った環境は次の通り。

$ sw_vers 
ProductName:    Mac OS X
ProductVersion: 10.13.4
BuildVersion:   17E199
$ python -V
Python 3.6.5

下準備

ひとまず pandas` をインストールしておく。

$ pip install pandas
$ pip list --format=columns | grep pandas
pandas          0.22.0

Python の REPL を起動する。

$ python

そして、ユーザの購買履歴っぽいサンプルデータを DataFrame オブジェクトで用意しておく。

>>> import pandas as pd
>>> columns = ['name', 'purchase_price']
>>> data = [
...     ('alice', 2000),
...     ('alice', 3000),
...     ('alice', 4000),
...     ('alice', 5000),
...     ('bob', 1000),
...     ('bob', 2000),
...     ('bob', 3000),
...     ('bob', 4000),
... ]
>>> df = pd.DataFrame(data, columns=columns)

これで下準備が終わった。

DataFrameGroupBy について

最初に、今回の主役となるオブジェクト DataFrameGroupBy について。

まず、先ほど作成した DataFrame オブジェクトに対して groupby() メソッドを使う。これによって、メソッドで指定したカラム単位で値を集約できる。

>>> dfg = df.groupby('name')

得られるオブジェクトは次の通り DataFrameGroupBy というものになる。

>>> dfg
<pandas.core.groupby.DataFrameGroupBy object at 0x103637780>

このオブジェクトには max() や min(), mean(), sum() といった代表的な統計量を計算するためのメソッドが用意されている。次の通り、集約した値の単位で統計量が計算される。

>>> dfg.max()
       purchase_price
name                 
alice            5000
bob              4000
>>> dfg.min()
       purchase_price
name                 
alice            2000
bob              1000
>>> dfg.sum()
       purchase_price
name                 
alice           14000
bob             10000
>>> dfg.mean()
       purchase_price
name                 
alice            3500
bob              2500

ただまあ、それぞれを単独で呼び出すのは結構めんどくさい。

そこで、代わりに agg() メソッドが使える。このメソッドは、辞書型のオブジェクトを渡すことでカラムに対して特定の集計をするように指示できる。それも、次のように値をリストにしておけば複数の集計が一度にできる。

>>> dfg.agg({
...     'purchase_price': ['max', 'min', 'sum', 'mean'],
... })
      purchase_price                   
                 max   min    sum  mean
name                                   
alice           5000  2000  14000  3500
bob             4000  1000  10000  2500

まあ、ただ上記の集計だけに関していえば describe() メソッドを使った方が楽かも。四分位数まで含めた代表的な統計量を一通り出力してくれる。

>>> dfg.describe()
      purchase_price                                                       \
               count    mean          std     min     25%     50%     75%   
name                                                                        
alice            4.0  3500.0  1290.994449  2000.0  2750.0  3500.0  4250.0   
bob              4.0  2500.0  1290.994449  1000.0  1750.0  2500.0  3250.0   

               
          max  
name           
alice  5000.0  
bob    4000.0

とはいえ agg() の本領は関数を渡せるところにあるんだと思う。具体的には、一つの引数を受け取る関数を辞書の値として渡すことができる。こうすると、辞書のキーに指定したカラムが Series オブジェクトとして関数に渡される。

例えば次のコードでは、引数の Series をソートした上で最大値と最小値を除いた合計値を計算している。まあ、サンプルの実用性は別として agg() に関数を渡すと、このように柔軟な集計が可能となる。

>>> dfg.agg({'purchase_price': lambda s: sum(sorted(s)[1: -1])})
       purchase_price
name                 
alice            7000
bob              5000

まあ Python 的には lambda を使うと可読性が犠牲になるので、ちゃんと関数を定義した方が好ましいかな。

>>> def sum_middle(series):
...     """最大値と最小値の要素を除いた合計を返す"""
...     sorted_series = sorted(series)
...     # 最大値と最小値を取り除く
...     middle_elements = sorted_series[1: -1]
...     # 合計値を返す
...     return sum(middle_elements)
... 
>>> dfg.agg({'purchase_price': sum_middle})
       purchase_price
name                 
alice            7000
bob              5000