pandasで他のカラムを条件に、新しいカラムを追加する
pandas
Lastmod: 2021-02-10

概要

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の代わりにすることが個人的な発見でした。

参考文献