one-hot encodingされた特徴量を逆に元のカテゴリー値に戻す

MovieLensのデータセットがこんな感じで処理がめんどくさかったのでメモとして

やりたいこと

以下のように one-hot encoding された状態で渡されたデータセットを

   movie_id  action  horror  romance  sf
0         1       1       0        0   0
1         2       0       0        1   0
2         2       1       0        0   0
3         3       0       0        0   1
4         3       1       0        0   0
5         4       0       1        0   0
6         5       0       0        0   1
7         5       0       1        0   0
8         5       1       0        0   0

以下のように one-hot encoding される前の categorical な状態に戻したい

   movie_id    genre
0         1   action
1         2  romance
2         2   action
3         3       sf
4         3   action
5         4   horror
6         5       sf
7         5   horror
8         5   action

やり方

以下のような関数を用意する

def convert_onehot_to_category(df, id_col, one_hot_columns, category_col='category'):
    df_concat = pd.DataFrame(columns=[id_col, category_col])
    for col in one_hot_columns:
        # 値が1以上のもののみ残す
        df_each = df[df[col] >= 1][[id_col, col]]
        # 値をカテゴリー値に置き換える
        df_each[col] = col

        df_each.columns = [id_col, category_col]
        df_concat = pd.concat([df_concat, df_each], axis=0)

    # 重複削除
    df_concat = df_concat.drop_duplicates().reset_index(drop=True).sort_values(by=id_col)
    return df_concat

以下のように、

• one-hot encoding 後のカラム名
• idの含まれるカラム
• カテゴリー値に変換後のカラム名

を渡せば

genres = ['action', 'romance', 'sf', 'horror']
id_col = 'movie_id'
category_col = 'genre'

df_category = convert_onehot_to_category(df_onehot, id_col=id_col, one_hot_columns=genres, category_col=category_col)

print(df_category)

元のカテゴリー値に変換してくれる

  movie_id    genre
0        1   action
1        2   action
4        2  romance
2        3   action
5        3       sf
7        4   horror
3        5   action
6        5       sf
8        5   horror