け日記

最近はPythonでいろいろやってます

小ネタ: Pandasのapplyで複数の列を返す

Pandasの小ネタです。

Pandasのapplyで1つの列から、複数の列を返したい場合があります。
その場合、1つ1つapplyで取り出しても良いのですが、冗長になりがちです。

import pandas as pd

df = pd.DataFrame({
    'number': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]
})

df['fizz'] = df['number'].apply(lambda n: n % 3 == 0)
df['buzz'] = df['number'].apply(lambda n: n % 5 == 0)
df['fizzbuzz'] = df['number'].apply(lambda n: n % 15 == 0)

以下のように、Seriesにまとめて返すと良いです。そのままDataFrameで列に設定することもできます。

df[['fizz', 'buzz', 'fizzbuzz']] = df['number'].apply(
    lambda n: pd.Series([n % 3 == 0, n % 5 == 0, n % 15 == 0])
)