Pandasのプログラムで、あるDataFrameから、関連する別のDataFrameに存在する値を含む条件に合う行を選択するコードの書き方を見つけるのに、結構時間が掛かった。
In [1]:
df1 = pd.DataFrame({
'name': list('AAABBBCCC'),
'val': range(1, 10)
})
df1
| name | val | |
|---|---|---|
| 0 | A | 1 |
| 1 | A | 2 |
| 2 | A | 3 |
| 3 | B | 4 |
| 4 | B | 5 |
| 5 | B | 6 |
| 6 | C | 7 |
| 7 | C | 8 |
| 8 | C | 9 |
こういうDataFrameと
In [2]:
df2 = df1.groupby('name').sum() + 1
df2.columns = ['val2']
df2
| val2 | |
|---|---|
| name | |
| A | 7 |
| B | 16 |
| C | 25 |
こういうDataFrameがあり、df1から"val2"の値が10より大きな行を選択したい、但しdf1とdf2とのmergeはワークメモリの都合でやりたくないとする。
df1とdf2をmergeして良いなら、やりたいことは次のことである。
In [3]:
df = df1.merge(df2, on='name')
df[df['val2'] > 10]
| name | val | val2 | |
|---|---|---|---|
| 3 | B | 4 | 16 |
| 4 | B | 5 | 16 |
| 5 | B | 6 | 16 |
| 6 | C | 7 | 25 |
| 7 | C | 8 | 25 |
| 8 | C | 9 | 25 |
次のように書ければ筆者としては直観的なのだが、これはエラーになる。
In [4]:
df1[df2.loc[df1['name'], 'val2'] > 10]
ValueError: cannot reindex from a duplicate axis
括弧内の"df2.loc[df1['name'], 'val2'] > 10"はdf1と同じ行数のboolean値を返すので、このままboolean indexingとして通してくれても良さそうなものである。
In [5]:
df2.loc[df1['name'], 'val2'] > 10
name A False A False A False B True B True B True C True C True C True Name: val2, dtype: bool
indexが問題とのことなので、次のようにすれば成功するが、無駄に複雑というか、そんな問題をいちいち意識したくない。
In [6]:
# 成功例1
df1[(df2.loc[df1['name'], 'val2'] > 10).reset_index(drop=True)]
| name | val | |
|---|---|---|
| 3 | B | 4 |
| 4 | B | 5 |
| 5 | B | 6 |
| 6 | C | 7 |
| 7 | C | 8 |
| 8 | C | 9 |
結局、Webで調べて試行錯誤しながら辿り着いたコードは、次のようにmapやapplyを使う形になった。
In [7]:
# 成功例2
df1[df1['name'].map(lambda x: df2.loc[x, 'val2'] > 10)]
| name | val | |
|---|---|---|
| 3 | B | 4 |
| 4 | B | 5 |
| 5 | B | 6 |
| 6 | C | 7 |
| 7 | C | 8 |
| 8 | C | 9 |
# 成功例3
df1[df1.apply(lambda df: df2.loc[df['name'], 'val2'] > 10, axis=1)]
（上と同じなので省略）
処理時間は、筆者の環境では成功例1と成功例3が大差なく、成功例2が半分ほどだった。
筆者が実際に困ったケースでは、2つのDataFrameを結合するキー（上の例の"name"）が複数の列からなっていた。
In [9]:
df3 = pd.DataFrame({
'name1': list('AAAABBBB'),
'name2': list('xxyyxxyy'),
'val': range(1, 9),
})
df3
| name1 | name2 | val | |
|---|---|---|---|
| 0 | A | x | 1 |
| 1 | A | x | 2 |
| 2 | A | y | 3 |
| 3 | A | y | 4 |
| 4 | B | x | 5 |
| 5 | B | x | 6 |
| 6 | B | y | 7 |
| 7 | B | y | 8 |
df4 = df3.groupby(['name1', 'name2']).agg(sum) + 1
df4.columns = ['val2']
df4
| val2 | ||
|---|---|---|
| name1 | name2 | |
| A | x | 4 |
| y | 8 | |
| B | x | 12 |
| y | 16 |
こういうDataFrameがあり、df3から"val2"の値が5より大きな行を選択したい、但しdf3とdf4とのmergeはワークメモリの都合でやりたくない、というような問題である。
DataFrameにはmapメソッドが無いので、成功例2の延長では書けず、成功例3の延長で、MultiIndexのindexingに苦戦しながら、次のようなコードに行き着いた。
# 成功例3'
df3[df3.apply(lambda df: df4.loc[tuple(df[['name1', 'name2']]), 'val2'] > 5, axis=1)]
| name1 | name2 | val | |
|---|---|---|---|
| 2 | A | y | 3 |
| 3 | A | y | 4 |
| 4 | B | x | 5 |
| 5 | B | x | 6 |
| 6 | B | y | 7 |
| 7 | B | y | 8 |
しかし、ここまで複雑になるのなら、成功列1のような方法でも良いかなとも思えてきた。
In [12]:
# 成功列1'
df3[(df4.loc[df3[['name1', 'name2']].values.tolist(), 'val2'] > 5).reset_index(drop=True)]
（上と同じなので省略）
筆者にはやはり成功例3'より美しくない上に一回りややこしく感じるが、筆者の環境では、成功例1'の処理時間は成功例3'の6割ほどだった。
コメント