DataFrameに「既存列を使った計算式の結果」を返す列を新たに作る方法

Question

既存データフレームdfに対して、条件式の結果適用される計算結果を新たな列に代入していきたいのですが・・・

●既存データフレーム
df
A B cls
15 25 1
25 10 2
35 40 3
40 55 5

●やりたいこと
pl列を新たに作り、clsが3より大きければ(B-A)*100、clsが3なら0、clsが2以下なら(A-B)*100を適用する
df
A B cls pl
15 25 1 -1000
25 10 2 1500
35 40 3 0
40 55 5 1500

過去似たような質問がありましたが、条件が単純だったこともありlambdaで記述されていました。できれば条件分岐の関数を新たに定義してapplyしたかったのですが、今回の場合への応用方法が分かりませんでした。

https://teratail.com/questions/27787

毎度お手数をおかけ致しますが、よろしくお願いいたします。

Answer 1

こんな感じでどうでしょうか？

import pandas as pd

df = pd.DataFrame({'A':[15,25,35,40],'B':[25,10,40,55], 'cls':[1,2,3,5]})

def f(row):
    if row.cls > 3:
        return (row.B - row.A) * 100
    elif row.cls == 3:
        return 0
    else:
        return (row.A - row.B) * 100

df['pl'] = df.apply(f, axis=1)
print(df)
#     A   B  cls    pl
# 0  15  25    1 -1000
# 1  25  10    2  1500
# 2  35  40    3     0
# 3  40  55    5  1500

Comments

[jun.k]

applyの前はdfで大丈夫なんですね。勉強になりました。いつも素早いご回答ありがとうございます!!

[magichan]

Series.apply() の場合は、引数としてセルの値が渡されてしまいますので、他のColumnの値を条件文で使用することができなくなります。
ですので、DataFrame.appley(axis=1) として引数に行データが渡るようにするのがよいかと思います。

[jun.k]

了解です。ありがとうございました。