python

python/pandasでピボットテーブルを自作する

pandasのpivot_tableがどうも気に入らないので書き直した。

まずかなり適当なサンプルデータ

import pandas as pd
import numpy as np

N=3000
d=pd.DataFrame({
    'job':np.random.randint(low=1, high=4, size=N),
    'age':np.random.randint(low=1, high=6, size=N),
})
d['v1'] = d['job']+(d['age'] >= 3).astype(int)*10+(d['age'] >= 5).astype(int)*10+np.random.normal(size=N)
d['v2'] = d['age']+d['job']-np.random.normal(size=N)
d['v3'] = (d['age']+d['job'])*0.1+np.random.normal(size=N)
d.loc[d['v3']>1, 'v3'] = 1
d.loc[d['v3']<0, 'v3'] = 0

作った関数に入れる。
集計したい変数名のリストをvarlistに、それに対する集計方法をmethodlistに順番に入れる。
mean, mで平均、sum, sで合計、ratio, rで平均の割合を見る。

new_pivot(
    d = d, v1='job', v2='age', 
    varlist = ['v1', 'v2', 'v3'],
    methodlist = ['m', 's', 'r']
)

すると次の表が出てくる。
カテゴリカル変数v1,v2の2軸に従って、指定した複数の変数それぞれを指定した方法で集計する。
あとはNとして各グループのサンプルサイズとデータ全体に対する割合も出す。

以下関数の中身。
なるべく再帰的に書きましたが、まだtotal出すために再集計するところを書き直したり、reset_index()しまくったりしているあたり美しくないですね。

def new_pivot(d, v1, v2, varlist, methodlist):
    N = len(d)
    d['value'] = 1
    
    for v, m in zip(['value']+varlist, ['sum']+methodlist):      
        if m in ['ratio', 'r', 'mean', 'm']:
            arg = np.mean
        elif m in ['sum', 's']:
            arg = np.sum
        
        t = d.groupby([v1,v2]).apply(arg)[[v]].reset_index(drop=False)

        _t = d.groupby([v1]).apply(arg)[[v]].reset_index(drop=False)
        _t[v2] = 'total'
        t = pd.concat([t,_t], axis=0, sort=False)

        _t = d.groupby([v2]).apply(arg)[[v]].reset_index(drop=False)
        _t[v1] = 'total'
        t = pd.concat([t,_t], axis=0, sort=False)

        _t = d.apply(arg)[[v]].to_frame().T
        _t[v1] = _t[v2] = 'total'
        t = pd.concat([t, _t], axis=0, sort=False)

        if m in ['ratio', 'r']:
            t[v] = (np.round(t[v], 3)*100).astype(str)+'%'
        elif m in ['mean', 'm']:
            t[v] = np.round(t[v], 1).astype(str)
        elif m in ['sum', 's']:
            total = d.sum()[[v]].values[0]
            t[v] = t[v].astype(int).astype(str)+' ('+(np.round(t[v]/total,3)*100).astype(str)+'%)'
        else:
            raise ValueError('Specify valid method')
            
        t = t.rename(columns={v:'value'})
        if 'tab' in locals():
            t['type'] = v+' ('+m+')'
            tab = pd.concat([tab, t], axis = 0, sort=False)
        else:
            t['type'] = 'N'
            tab = t
    
    tab = tab.set_index([v1, v2, 'type']).unstack(v2).fillna(0)
    display(tab)

これベースにいくらでも拡張できますね。
集計用の変数v1/v2が連続値だった場合には四分位点に従ったカテゴリカルに変換してから集計するとか。

まあこんなんでよければ好きに使ってください。