如何改善此代码以使用OneHotEncoder?

时间:2019-11-20 19:18:19

标签: python pandas scikit-learn

您好,我正在使用以下函数将数据集中的所有类别值转换为数值,但是我想将其转换为使用OneHotEncoder,该怎么做?

def categorical_to_numerical(dataframe):
    for col in dataframe.columns:
        if str(dataframe[col].dtype) == 'category':
            dataframe[col] = dataframe[col].astype("category").cat.codes
    return dataframe

谢谢

1 个答案:

答案 0 :(得分:1)

如果我对您的理解正确,那么您想使用DataFrame.select_dtypes来选择object (string)列。

# example dataframe
df = pd.DataFrame({'col1':[1,2,3],
                   'col2':['a','b','a'],
                   'col3':[4,5,6],
                   'col4':['aaa', 'bbb', 'bbb']})

   col1 col2  col3 col4
0     1    a     4  aaa
1     2    b     5  bbb
2     3    a     6  bbb
for col in df.select_dtypes('object'):
    df[col] = df[col].astype('category').cat.codes

   col1  col2  col3  col4
0     1     0     4     0
1     2     1     5     1
2     3     0     6     1

或者,如果您想实际使用OneHotEncode,我们可以使用pd.get_dummies

df = pd.get_dummies(df)

   col1  col3  col2_a  col2_b  col4_aaa  col4_bbb
0     1     4       1       0         1         0
1     2     5       0       1         0         1
2     3     6       1       0         0         1