2015年12月05日

pandasのapplyで複数の値を同時に返したい

こんなDataFrameがあったとする。

import pandas as pd
df = pd.DataFrame( [{ '店舗': '渋谷支店', '価格': 500 }, { '店舗': '新宿支店', '価格': 515 }] )

df
  #=>     価格    店舗
  #=> 0  500  渋谷支店
  #=> 1  515  新宿支店

ここに消費税額と税込み金額を加えたいとする。

df['税込価格'] = (df['価格'] * 1.08).astype(np.int)
df['税額'] = df['税込価格'] - df['価格']
df

  #=>     価格    店舗  税込価格  税額
  #=>  0  500  渋谷支店   540  40
  #=>  1  515  新宿支店   556  41

個別にapplyで実行するならこんな感じ。

def tax1( row ):
    return int(row['価格'] * 1.08)

def tax2( row ):
    return row['税込価格'] - row['価格']

df['税込価格'] = df.apply(tax1, axis=1)
df['税額'] = df.apply(tax2, axis=1)
df

  #=>     価格    店舗  税込価格  税額
  #=>  0  500  渋谷支店   540  40
  #=>  1  515  新宿支店   556  41

これがもっと複雑な計算で、可能であれば1回のapplyで2つの結果を返すようにしたかっとする。その場合はSeriesを返してaxis=1でconcatすれば良いらしい。

def tax( row ):
    x = int(row['価格'] * 1.08)
    y = x - row['価格']
    return pd.Series({'税込価格': x, '税額': y})

pd.concat( [df, df.apply( tax, axis=1 )], axis=1 )

  #=>     価格    店舗  税込価格  税額
  #=>  0  500  渋谷支店   540  40
  #=>  1  515  新宿支店   556  41

rowレベルでapplyすると重いので上記のような例では個別に計算した方が良いけど、両者の計算が絡み合っていて複数同時に値を戻した方が効率が良い場合はこれを使いそう。