pandasでリストとして格納されている各要素をカラムとして設定し、ワンホットエンコードとして展開する

データ分析をしていて、pandasの要素にリストが格納されており、そのリストに対してワンホットエンコードした状態のDataFrameを作りたいという機会があり、結構苦労したのでメモしておきます。

github

  • jupyter notebook形式のファイルはこちら

google colaboratory

  • google colaboratory で実行する場合はこちら

実行環境

筆者のOSはmacOSです。LinuxやUnixのコマンドとはオプションが異なります。

!sw_vers

ProductName:		macOS
ProductVersion:		13.5.1
BuildVersion:		22G90

!python -V

Python 3.9.17

基本的なライブラリをインポートしそのバージョンを確認しておきます。

%matplotlib inline

import pandas as pd

print('pandas version :', pd.__version__)

pandas version : 2.0.3

サンプルデータの準備

df = pd.DataFrame(
    {
        "user_id": ["A", "B", "C"],
        "item_id": [["PC", "Book", "Water"], ["Book", "Table"], ["Desk", "CD"]],
    }
)

df.head()
user_iditem_id
0A[PC, Book, Water]
1B[Book, Table]
2C[Desk, CD]

MultiLabelBinarizerの利用

結論から述べると、MultiLabelBinarizer というscikit-learnのライブラリを利用します。

以下のように、 fit_transformを利用する事で、ワンホットエンコードを簡単に実現できます。また、それに対応するカラム名も簡単に取得できます。

from sklearn.preprocessing import MultiLabelBinarizer


mlb = MultiLabelBinarizer()
mlb.fit_transform(df.item_id)

array([[1, 0, 0, 1, 0, 1],
       [1, 0, 0, 0, 1, 0],
       [0, 1, 1, 0, 0, 0]])

mlb.classes_

array(['Book', 'CD', 'Desk', 'PC', 'Table', 'Water'], dtype=object)

あとはこれを組み合わせるだけです。popでdfから取りだして、最後にjoinで結合します。

out_df = df.join(pd.DataFrame(mlb.fit_transform(df.pop("item_id")), columns=mlb.classes_))

out_df
user_idBookCDDeskPCTableWater
0A100101
1B100010
2C011000

参考サイト