pandasで値がTrueの時に、対応する列ラベルを取り出したい。短時間で実行するにはどうしたらよいか。

Question

###前提・実現したいこと
以下のようなブーリアン型の表があるときに、Trueの値がある列のラベルを取り出してまとめたいです。

処理前の表
|product|key1|key2|key3|
|:--|:--:|--:|
|ProductA|True|False|False|
|ProductB|False|Ture|False|
|ProductC|True|False|True|

処理後の表
|product|key1|key2|key3|keys|
|:--|:--:|--:|--:|
|ProductA|True|False|False|key1|
|ProductB|False|Ture|False|key2|
|ProductC|True|False|True|key1,key3|

###発生している問題・エラーメッセージ
以下のようなfor文を使ったコードで目的にあったリストを作ることができたのですが、大きなデータフレーム(400列、15万行)に適用すると実行に数時間かかってしまいます。もっと短時間で実行できるような書き方はできないでしょうか？

###該当のソースコード

import pandas as pd
import numpy as np

df = pd.DataFrame([
  ["ProductA", True, False, False],
  ["ProductB", False, True, False],
  ["ProductC", True, False, True]],
  columns=['product', 'key1', 'key2', 'key3'])

keys = []
for col in range(0,len(df.index)):
    key = []
    for row in range(1,len(df.columns)):
        if df.iloc[[col],[row]].values:
            key.append(df.columns[row])
    keys.append(key)
keys

Answer 1

pandasで書くならこんな感じですかね。

import pandas as pd

# 行ごとのデータを受け取って、Trueの列だけのindexをリストに変換してjoinするlambda関数を定義
f = lambda series: ",".join(series.index[series==True].tolist())

# 入出力例
f(df.iloc[2])
# >'key1,key3'

# applyでaxis=1（行方向に演算）でfuncを適用する
# 結果をdfのkeys列に新規に挿入
df['keys'] = df.apply(f, axis = 1)

df
# >     product   key1   key2   key3       keys
# > 0  ProductA   True  False  False       key1
# > 1  ProductB  False   True  False       key2
# > 2  ProductC   True  False   True  key1,key3

無理にlambdaを使わなくても普通の関数でもOKです。

def func(series):
    columnNames=series.index[series==True].tolist()
    return ",".join(columnNames)

df['keys'] = df.apply(func, axis = 1)

Comments

[tecton00]

ご教示ありがとうございます。
おかげさまで実際のデータでも短時間で実行することができました。
lambda関数は今まで避けてきたのですが、きちんと勉強しようと思います。

Answer 2

以下が時間消費の主因になっていると考えます。
0. forループを二重に回しているところ。
0. ループの最奥でappendを一々実行しているところ。

これを解消するため以下の点に留意してスクリプト後半を書きかえてみました。
0. ループの外に出せる処理は全て外に出す。
0. Trueである列の列番号を行毎に拾う処理をリスト内包表記で記述して、forループを一つ消す。
0. 拾った列番号は別リスト(df1)に格納。それ以外の処理はループ内では一切行わない。

df1  = list(np.repeat("",len(df.index)))
l    = len(df.ix[0])
cols = list(df)
 
for row in range(0,len(df.index)):
    df1[row] = [x[0] for x in enumerate(df.ix[row, range(1,l)]) if x[1]]

# とりあえず第2行目の結果をそれっぽく出力してみる。
print df.ix[1,0]+"\t", "\t".join( [cols[x+1] for x in df1[1]] )

5000行 x 400列のデータを処理してみたところ、質問文にあるスクリプトでは 1280秒掛かった処理が、上記改変版では7秒で終わりました。

Comments

[tecton00]

ご教示ありがとうございました！
私の環境でも数時間かかっていた処理が数十秒で実行できました。
考え方も教えていただきありがとうございます。データ量が多くなるとコードの書き方ひとつで処理能力が大きく変わることを実感できました。