TL;DR
-
pandas.DataFrame.applyのオプショナル引数に、result_type="expand"を指定する -
pandas.DataFrame.applyの引数の関数(ラムダ式)は、タプルまたはリストを返すようにする - 代入式の左辺では、追加する列名をリストで指定する
def get_values(value0):
# some calculation
return value1, value2
df[["column1", "column2"]] = df.apply(
lambda r: get_values(r["column0"]),
axis=1,
result_type="expand")
解説
適当なpandas.DataFrameがあるとします。
import pandas as pd
df = pd.DataFrame()
df["x"] = [1, 2, 3, 4, 5]
print(df)
x
0 1
1 2
2 3
3 4
4 5
このxという名前の列を使って計算を行い、新たな列を作成したいとします。ベクトルのまま計算できる場合は簡単です。ここでは2乗を計算する場合を例に取ってみます。
df["x2"] = df["x"] ** 2
print(df)
x x2
0 1 1
1 2 4
2 3 9
3 4 16
4 5 25
もし新たな列の求める計算が複雑な関数で定義されており、ベクトルのまま計算できない場合には、pandas.DataFrame.applyを用いることができます。
def square(x):
return x ** 2
df["x2"] = df.apply(lambda r: square(r["x"]), axis=1)
print(df)
もちろん結果は前と同じです。
さて、以下のような複数の値を返す関数を使って、複数列を同時に追加する場合はどうなるでしょうか？
def get_values(x):
return x ** 2, x ** 3
次のように書くとエラーになります。
df[["x2","x3"]] = df.apply(lambda r: get_values(r["x"]), axis=1)
---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
<ipython-input-19-3039b5c1e701> in <module>
1 def get_values(x):
2 return x**2, x**3
----> 3 df[["x2", "x3"]] = df.apply(lambda r: get_values(r["x"]), axis=1)
(中略)
ValueError: Must have equal len keys and value when setting with an iterable
このエラーが起きている理由は、右辺の値をprintすることで分かります。
print(df.apply(lambda r: get_values(r["x"]), axis=1))
0 (1.0, 1.0)
1 (4.0, 8.0)
2 (9.0, 27.0)
3 (16.0, 64.0)
4 (25.0, 125.0)
dtype: object
2個の要素を持つタプルが入った(1列の)Seriesになってしまっています。
2列5行のDataFrameにするには、result_type="expand"を指定します。
df[["x2","x3"]] = df.apply(lambda r: get_values(r["x"]), axis=1, result_type="expand")
print(df)
x x2 x3
0 1 1.0 1.0
1 2 4.0 8.0
2 3 9.0 27.0
3 4 16.0 64.0
4 5 25.0 125.0