(Pandas)複数の列を使ってtransformする方法で悩む

In [1]:
import pandas as pd
df = pd.DataFrame({'name': ['a', 'b', 'c', 'd'] * 3,
                               'value': range(12),
                               'flag': [1, 0, 0, 0, 1, 0] * 2})
df
Out [1]:
name value flag
0 a 0 1
1 b 1 0
2 c 2 0
3 d 3 0
4 a 4 1
5 b 5 0
6 c 6 1
7 d 7 0
8 a 8 0
9 b 9 0
10 c 10 1
11 d 11 0

こういうDataFrameがあり、'name'でgroupbyして、'flag'に1つでも1があるグループの、'flag'が1の行がグループ内の'value'の平均、それ以外の行は0という列(この例ではインデックスが0と4の所が'a'の平均、6と10の所が'c'の平均、それ以外は0という列)を追加したいとする。
大体そういう感じのことをtransformでやりたかったが、スマートなやり方がわからず、悩んでいる。

目的の列をグループ毎に返す関数を、transformの代わりにapplyを使って、

In [2]:

def func(groupdf):
    ret = pd.Series(0, index=groupdf.index)
    if any(groupdf['flag']):
        ret[groupdf['flag'] == True] = groupdf['value'].mean()
    return ret
    
df.groupby('name').apply(func)

又は

In [3]:

df.groupby('name').apply(lambda x: x['value'].mean() * x['flag'] * any(x['flag']))

Out [3]:

name    
a     0     4.0
      4     4.0
      8     0.0
b     1     0.0
      5     0.0
      9     0.0
c     2     0.0
      6     6.0
      10    6.0
d     3     0.0
      7     0.0
      11    0.0
Name: flag, dtype: float64

というように作ることができたのだが、[2]のfuncも[3]のlambda関数も、transformに与えると、'flag'という列が無いというエラーになる。
applyなら呼び出される関数に複数列のDataFrameが渡されるので複数列を参照しながら計算ができるが、aggregateやtransformだと呼び出される関数に1列分のSeriesしか渡されないので、複数列を参照しながら計算ができない。

もし、flagが1の行だけがグループの平均という条件を外し、'flag'に1つでも1があるグループはグループ内の'value'の平均、それ以外の行は0という列で良い、つまりグループ内は全て同じ値になるなら、Webでサンプルコードがいくつか見つかり、大きく分けて2つの方法が見つかった。
1つは、列毎にtransformした結果を組み合わせて何とかするという方法である。

In [4]:

grouped = df.groupby('name')
df['ave'] = grouped['value'].transform(np.mean) * grouped['flag'].transform(any)
df
Out [4]:
(省略)

もう1つは、applyを使ってグループ毎に計算した結果をmergeする方法である。

In [5]:

_ = df.groupby('name').apply(lambda x: x['value'].mean() if any(x['flag']) else 0).rename('ave')
df = df.merge(_, how='left', on='name')
df
Out [5]:
name value flag ave
0 a 0 1 4
1 b 1 0 0
2 c 2 0 6
3 d 3 0 0
4 a 4 1 4
5 b 5 0 0
6 c 6 1 6
7 d 7 0 0
8 a 8 0 4
9 b 9 0 0
10 c 10 1 6
11 d 11 0 0

前者は列毎にtransformした後で何とかなる場合しか使えないのに対し、後者は複数の列を参照しながら計算できるので、後者の方が汎用的だと思う。
速度面では、前者は中間データを作成して処理時間がかかりがちなtransformを複数回実行するので不利なように思えたが、筆者のJupyter Notebookの%timeitで計測した限り、DataFrameのサイズを10,000倍とかにしても、処理時間は大差なかった。

元のやりたいことについては、前者(In [4]の例)を応用すると、次のようなのができた。

In [6]:

grouped = df.groupby('name')
df['ave'] = grouped['value'].transform(np.mean) * grouped['flag'].transform(any) * df['flag']
df
Out [6]:
name value flag ave
0 a 0 1 4.0
1 0 0 0 0.0
2 0 0 0 0.0
3 0 0 0 0.0
4 a 4 1 4.0
5 0 0 0 0.0
6 c 6 1 6.0
7 0 0 0 0.0
8 0 0 0 0.0
9 0 0 0 0.0
10 c 10 1 6.0
11 0 0 0 0.0

しかし、transformを2回使うし、式が場当たり的で汎用的でないのが不満である。
後者(In [5]の例)を応用すると、次のようにするしか思い付かない。

In [7]:

_ = df.groupby('name').apply(lambda x: x['value'].mean() if any(x['flag']) else 0).rename('ave')
df = df.merge(_, how='left', on='name')
df[df['flag'] == 0] = 0
df
Out [7]:
(Out [6] と同じ)

これも、後のboolean indexing部分はグループを無視して処理しているので、そのようにできない時は同じようにできないし、これによって処理時間が大幅に伸びるし、transformの出番と思われるのにtransformを使ってないのが不満である。


Pandasのgroupbyとは3年の付き合いなのだが、たまにしか触らないせいか、仕様が複雑なせいか、なかなか頭に定着しない。こういう感じのコードにすれば良さそう、と思って書き始めても、半分くらいは全く違うコードになる。最初にコードのイメージが全くわからない時もある。何をした時に、行インデックスと列インデックスが何になって、データがどうなるかを、きちんと理解できてないからなんだろうか。