Edited at

pandasで条件分岐(case when的な)によるデータ加工を網羅したい

More than 1 year has passed since last update.

pandasでデータ加工をしつつ分析を進めていると、大抵の処理は問題なくできるのだけれど、SQLのcase whenのような条件分岐の処理をpandasでやろうとするとあまりの面倒くささに死にたくなることがよくある。

しかも毎回やり方を忘れるからここに書いておくことにする。

できるだけSQLと対比して書くことにする。

今回は2017年の日経平均株価を例にcase whenのパターンを書いてみたい。

元データはこんな感じ。

dt
stock_id
open
high
low
close
market_id

2017-01-04
1001
19298
19594
19277
19594
1

2017-01-05
1001
19602
19615
19473
19520
1

2017-01-06
1001
19393
19472
19354
19454
1

2017-01-10
1001
19414
19484
19255
19301
1

2017-01-11
1001
19358
19402
19325
19364
1


  • dt:日付

  • stock_id:株価ID

  • open:始値

  • high:最高値

  • low:最低値

  • close:終値

  • market_id:市場ID


1変数が条件分岐の対象&バイナリ化したい場合

条件分岐の対象が1つだけの場合には、pandasで加工するにしてもSQLで加工するにしてもそれほど大変ではない。。

さらに、[0,1]のように2値の値を取るように変換する場合はさらにラク。

例えば、丁度日経平均が2万円台に到達しようとするあたりなので、openが2万円を超えたら1がつくようなflg変数を追加する場合を考える。

SQLでは簡単に書ける。

select

*,
case when open >= 20000 then 1 else 0 end as flg_20t --20000円以上
from
sample_data --仮にテーブル名はsample_dataとする

では、pandasではどうなるかと言うと、まだ1行で書ける。

sample_data['flg_open_20t'] = sample_data['open'].apply(lambda x : 1 if x >= 20000 else 0)

sample_data.loc[sample_data.dt>='2017-06-01',:].head()

dt
stock_id
open
high
low
close
market_id
flg_open_20t

2017-06-01
1001
19692
19887
19686
19860
1
0

2017-06-02
1001
19970
20239
19967
20177
1
0

2017-06-05
1001
20135
20224
20104
20170
1
1

2017-06-06
1001
20122
20152
19948
19979
1
1

2017-06-07
1001
19951
20023
19908
19984
1
0

うん、できてるね。

次。


1変数が条件分岐の対象&複数の状態に分岐させたい(カテゴリ化したい)

次は、case whenの条件分岐の対象になる列は1つのまま変わらないのだけれど、FLGではなくてカテゴリ変数のように複数の状態を表現する場合を考える。

例えば、

1. openが1万円未満→0

2. openが1万円以下ー1万5千円未満→1

3. openが1万5千円以下ー2万円未満→2

4. openが2万円以上→3

を表現したい。

SQLでは、こんな感じ。まだまだ楽勝ですわ。

select

*,
case when open < 10000 then 0
when open >= 10000 and open < 15000 then 1
when open >= 15000 and open < 20000 then 2
else 3 end as cate --カテゴリ変数
from
sample_data --仮にテーブル名はsample_dataとする

さて、pandasではどうかというと

sample_data['cate'] = sample_data['open'].apply(lambda x : 0 if x < 10000 else (1 if x >= 10000 and x < 15000 else(2 if x >= 15000 and x < 20000 else 3)))

だいぶ複雑。これならもうlambdaは使わずに普通に関数にした方が良さそう。

ということで

# 条件分岐の関数を定義

def func_cate(x):
if x < 10000:
return 0
elif x >= 10000 and x < 15000:
return 1
elif x >= 15000 and x < 20000:
return 2
else:
return 3

sample_data['cate_f'] = sample_data['open'].apply(func_cate)
sample_data.head()

dt
stock_id
open
high
low
close
market_id
flg_open_20t
cate
cate_f

2017-01-04
1001
19298
19594
19277
19594
1
0
2
2

2017-01-05
1001
19602
19615
19473
19520
1
0
2
2

2017-01-06
1001
19393
19472
19354
19454
1
0
2
2

2017-01-10
1001
19414
19484
19255
19301
1
0
2
2

2017-01-11
1001
19358
19402
19325
19364
1
0
2
2

でけた。

次。


1変数が条件分岐の対象&元の値を変換

上の2つは1つの変数を条件分岐の対象としつつ、新たに作成する変数の数値も外側から与えるタイプの処理だったので、まだなんとかなった。

でも、元の値の一部を置換するタイプの処理はだいぶ面倒になる。

例えば、こんな場合。

openの値が

1. 1万9000円未満ならばNULLに置換

2. 1万9000円以上2万円未満ならば19500に置換

3. 2万円以上ならopenの値そのまま

データに外れ値があって特定の値に置換して使いたい場合にやる必要があることが多い。

SQLなら別に大したことない。

select

*,
case when open < 19000 then NULL
when open >= 19000 and open < 20000 then 19500
else open end as open_2 --加工済みopen値
from
sample_data --仮にテーブル名はsample_dataとする

pandasなら


import numpy as np
def func_open_2(x):
if x < 19000:
return np.NaN
elif x >= 19000 and x < 20000:
return 19500
else:
return x

sample_data['open_2'] = sample_data['open'].apply(func_open_2)

dt
stock_id
open
high
low
close
market_id
flg_open_20t
cate
cate_f
open_2

2017-01-04
1001
19298
19594
19277
19594
1
0
2
2
19500.0

2017-01-05
1001
19602
19615
19473
19520
1
0
2
2
19500.0

2017-01-06
1001
19393
19472
19354
19454
1
0
2
2
19500.0

2017-01-06
1001
19393
19472
19354
19454
1
0
2
2
19500.0

2017-01-18
1001
18753
18941
18650
18894
1
0
2
2
NaN

2017-11-24
1001
22390
22567
22381
22550
1
1
3
3
22390.0

うん。ええやろ。

次。


複数変数が条件分岐の対象&カテゴリ化したい

ここからが本丸。

本当にpandasが私を苦しめ始めるのはこのあたりから。。。

1つの列だけを条件分岐の対象としている間は、列を特定してapplyを使っておけば良かったのだが、複数列が条件分岐に必要となると途端にpandasは扱いが面倒になる。

こういう場合を考えてみる。

1. openがcloseよりも低い(値上がり)→1

2. openがcloseよりも高い(値下がり)→2

3. openがcloseと同じ→0

SQLだったら簡単に実現できる。

select

*,
case when open < close then 1
when open > close then 2
else 0 end as cate_change --価格変動カテゴリ
from
sample_data --仮にテーブル名はsample_dataとする

pandasなら

def func_cate_change(x):

if x.open < x.close:
return 1
elif x.open > x.close:
return 2
else:
return 0

sample_data['cate_change'] = sample_data.apply(lambda x:func_cate_change(x),axis=1)

applyの処理がpandas.Seriesからpandas.DataFrameに変わっているのと処理の方向が列方向に対するものに変わっているのがポイント

このaxisの処理を忘れてよく何をやっているかわからなくなる orz

axis=0 →行方向(デフォルト)

axis=1 →列方向

dt
stock_id
open
high
low
close
market_id
flg_open_20t
cate
cate_f
open_2
cate_change

2017-01-04
1001
19298
19594
19277
19594
1
0
2
2
19500.0
1

2017-01-05
1001
19602
19615
19473
19520
1
0
2
2
19500.0
2

2017-01-06
1001
19393
19472
19354
19454
1
0
2
2
19500.0
1

2017-01-10
1001
19414
19484
19255
19301
1
0
2
2
19500.0
2

いいね!!

ちなみに

# 別にこっちでも通る

sample_data['cate_change_2'] = sample_data.apply(func_cate_change,axis=1)


複数変数が条件分岐の対象&値を加工する

条件が複数列にまたがるだけならまだしも、新たに作る列に用いる値が複数にまたがることもある。

例えば、

1. highとlowが100円以上離れているなら、openとcloseの平均を取る

2. highとlowが100円未満なら、highとlowの平均を取る

SQLなら

select

*,
case when high - low >= 100 then (open + close) /2
else (high + low) /2 end as avg_100 --変則平均値
from
sample_data --仮にテーブル名はsample_dataとする

pandasなら

def func_avg_100(x):

if x.high - x.low <= 100:
return (x.high + x.low) /2
else:
return (x.open + x.close) /2

sample_data['avg_100'] = sample_data.apply(lambda x:func_avg_100(x),axis=1)

dt
stock_id
open
high
low
close
market_id
flg_open_20t
cate
cate_f
open_2
cate_change
cate_change_2
avg_100

2017-01-04
1001
19298
19594
19277
19594
1
0
2
2
19500.0
1
1
19446.0

2017-01-05
1001
19602
19615
19473
19520
1
0
2
2
19500.0
2
2
19561.0

2017-01-06
1001
19393
19472
19354
19454
1
0
2
2
19500.0
1
1
19423.5

2017-01-10
1001
19414
19484
19255
19301
1
0
2
2
19500.0
2
2
19357.5

2017-01-11
1001
19358
19402
19325
19364
1
0
2
2
19500.0
1
1
19361.0

面倒だけど、まぁセーフ


複数変数が条件分岐の対象&外側から引数を与える

ここまでは、条件や値が複数になることはあっても、どのような条件で分岐するかやどのような値を取るかということは所与の値を決めて作っていた。

しかし、このままだと条件分岐のパターンを作るたびに関数を作成する必要があり非常に面倒。

そこで、引数を加えて任意の条件や値でも分岐できることを考える。

想定するパターンとしては、下記を想定。


  1. openがcloseに対して+X円以上ならYを代入

  2. openがcloseに対して+X円未満ならZを代入

SQLでは、プロシージャや関数を使わないと基本的にはかけないので飛ばす。

pandasなら

def func_xyz(row,x,y,z):

if row.open >= x:
return y
else:
return z

# x,y,zを事前に定義
x = 19500
y = 0
z = 1
sample_data['xyz'] = sample_data.apply(lambda row:func_xyz(row,x,y,z),axis=1)

引数を調整すれば、これで任意の変数を自由に条件分岐に使える

listとかdict使って制御した方が、頭良さそうだけど、まぁええやろ


複数変数が条件分岐の対象&欠損値置換

最後に欠損値の置換を条件に応じて変える場合を考える。


  1. 値下がりする日(cate_change:1)のopen_2が欠損値ならば、期間全体の平均値を置換する

  2. 値上がりする日(cate_change:2)のopen_2が欠損値ならば、当日のopenを置換する

  3. 変動がない日(cate_change:0)のopen_2が欠損値ならば、そのままopen_2を採用する

  4. 欠損値でないならそのままopen_2を採用する

まずは、SQLでかくと

select

a.*,
case when open_2 is NULL
case when cate_change = 1 then b.avg_open
when cate_change = 2 then a.open
else a.open_2 end
else open_2 end as open_3

from
sample_data as a,
--単純結合
(
select
avg(open) as avg_open
from
sample_data
) as b
;

SQLで書いてもかなり複雑だわ。。

pandasなら

import math

def func_open_3(row,avg_open):
# スカラーでnp.NaNを評価したい場合は、mathで判定するのがいいっぽい
if math.isnan(row.open_2):
if row.cate_change == 1:
return avg_open
elif row.cate_change == 2:
return row.open
else:
return row.open_2
else:
return row.open_3

# 行に対する処理と列に対する処理を同時に行うのは困難なので、まず欠損値処理を行う
sample_data['open_3'] = sample_data.open_2.fillna(method='bfill')

# 列に対する処理を実施する前に、open_2の行方向の平均値を事前に計算しておく
avg_open = sample_data['open'].mean()

# 列方向に対する処理を行う
sample_data['open_3'] = sample_data.apply(lambda x: func_open_3(x,avg_open),axis=1)

行方向の処理(平均や最大値の取得)と列方向の処理(条件分岐・置換)が込み入ってるとかなり複雑になってくるイメージ。

dt
stock_id
open
high
low
close
market_id
flg_open_20t
cate
cate_f
open_2
cate_change
cate_change_2
avg_100
xyz
open_3

3931
2017-01-17
1001
19038
19043
18812
18813
1
0
2
2
19500.0
2
2
18925.5
1

3932
2017-01-18
1001
18753
18941
18650
18894
1
0
2
2
NaN
1
1
18823.5
1

3933
2017-01-19
1001
19082
19122
18982
19072
1
0
2
2
19500.0
2
2
19077.0
1

3935
2017-01-23
1001
18938
19024
18879
18891
1
0
2
2
NaN
2
2
18914.5
1

ということで、色々やったから忘れた時、ここ見直せば今後自分を助けてくれるはず。。。

おしまい