概要
pandasでデータの前処理をしている時、現在あるカラムを条件に使って新しいカラムを作成したい時ってありますよね。今回はその方法を紹介します
方法
結論からいうとnp.selectを使うと柔軟に条件を指定することができます。
サンプルデータ準備
>>> import pandas as pd
>>> import numpy as np
>>>
>>> df = pd.DataFrame([["man", 81], ["man", 63], ["woman", 89], ["man", 93], ["woman", 56], ["woman", 76]], columns=['gender', 'score'])
>>>
>>> df
gender score
0 man 81
1 man 63
2 woman 89
3 man 93
4 woman 56
5 woman 76
1つの条件、分岐が1つの場合
条件の数、分岐の数に限らずconditionsに条件を書き、choicesに条件の結果を書き込んでいきます。
今回は80点以上取っている人に1, そうでないものに0を付与する想定をします。
>>> conditions = [
... df['score'] >=80
... ]
>>>
>>> choices = [1]
>>>
>>> df['test1'] = np.select(conditions, choices, default=0) #defaultを入れることでelseの代わりになる！
>>>
>>> df
gender score test1
0 man 81 1
1 man 63 0
2 woman 89 1
3 man 93 1
4 woman 56 0
5 woman 76 0
複数の条件、分岐が一つの場合
genderがmanで80点以上の場合は1, それ以外は0と想定。
>>> conditions = [
... (df['score'] >= 80) & (df['gender'] == 'man')
... ]
>>>
>>> choices = [1]
>>>
>>> df['test2'] = np.select(conditions, choices, default=0)
>>>
>>>
>>> df
gender score test1 test2
0 man 81 1 1
1 man 63 0 0
2 woman 89 1 0
3 man 93 1 1
4 woman 56 0 0
5 woman 76 0 0
複数の条件、分岐が複数の場合
genderがmanで80点以上の場合は2, genderがwomanで80点以上の場合は1, それ以外は0と想定。
>>> conditions = [
... (df['score'] >= 80) & (df['gender'] == 'man'),
... ((df['score'] >= 80) & (df['gender'] == 'woman'))
... ]
>>>
>>> choices = [2, 1]
>>>
>>> df['test3'] = np.select(conditions, choices, default=0)
>>>
>>>
>>> df
gender score test1 test2 test3
0 man 81 1 1 2
1 man 63 0 0 0
2 woman 89 1 0 1
3 man 93 1 1 2
4 woman 56 0 0 0
5 woman 76 0 0 0
まとめ
default設定でelseの代わりにすることが個人的な発見でした。