Apply one-hot encoding to a pandas DataFrame

one-hot.py

import pandas as pd
import numpy as np
from sklearn.feature_extraction import DictVectorizer
 
def encode_onehot(df, cols):
    """
    One-hot encoding is applied to columns specified in a pandas DataFrame.
    
    Modified from: https://gist.github.com/kljensen/5452382
    
    Details:
    
    http://en.wikipedia.org/wiki/One-hot
    http://scikit-learn.org/stable/modules/generated/sklearn.preprocessing.OneHotEncoder.html
    
    @param df pandas DataFrame
    @param cols a list of columns to encode
    @return a DataFrame with one-hot encoding
    """
    vec = DictVectorizer()
    
    vec_data = pd.DataFrame(vec.fit_transform(df[cols].to_dict(outtype='records')).toarray())
    vec_data.columns = vec.get_feature_names()
    vec_data.index = df.index
    
    df = df.drop(cols, axis=1)
    df = df.join(vec_data)
    return df


def main():
  np.random.seed(42)
  df = pd.DataFrame(np.random.randn(25, 3), columns=['a', 'b', 'c'])

  # Make some random categorical columns
  df['e'] = [random.choice(('Chicago', 'Boston', 'New York')) for i in range(df.shape[0])]
  df['f'] = [random.choice(('Chrome', 'Firefox', 'Opera', "Safari")) for i in range(df.shape[0])]

  # Vectorize the categorical columns: e & f
  df = encode_onehot(df, cols=['e', 'f'])
  print df.head()
 
if __name__ == '__main__':
    main()

One-hot encoding is supported in pandas (I think since 0.13.1) as pd.get_dummies.

@carlgieringer get_dummies() creates a new column for each value typical one-hot encoding just maps each value in a column to an integer.

@dreyco676, @carlgieringer is right, the following will transform a given column into one hot. Use prefix to have multiple dummies.

>>> import pandas as pd
>>> df = pd.DataFrame({'A': ['a', 'b', 'a'], 'B': ['b', 'a', 'c']})
>>> # Get one hot encoding of columns B
>>> one_hot = pd.get_dummies(df['B'])
>>> # Drop column B as it is now encoded
>>> df = df.drop('B', axis=1)
>>> # Join the encoded df
>>> df = df.join(one_hot)
>>> df
   A    a    b    c
0  a  0.0  1.0  0.0
1  b  1.0  0.0  0.0
2  a  0.0  0.0  1.0

An alternative approach that uses pd.factorize to map each item to a value in a single column is:

>>> import pandas as pd
>>> df = pd.DataFrame({'A': ['a', 'b', 'a','c'], 'B': ['a', 'b', 'a','c']})
>>> # Get encoding of column B
>>> catenc = pd.factorize(df['B'])
>>> catenc
(array([0, 1, 0, 2]), Index([u'a', u'b', u'c'], dtype='object'))
>>> # Add encoded column
>>> df['B_enc'] = catenc[0]
>>> df
   A  B  B_enc
0  a  a      0
1  b  b      1
2  a  a      0
3  c  c      2

...or skipping the intermediate step:

>>> import pandas as pd
>>> df = pd.DataFrame({'A': ['a', 'b', 'a','c'], 'B': ['a', 'b', 'a','c']})
>>> df['B_enc'] = pd.factorize(df['B'])[0]
>>> df
   A  B  B_enc
0  a  a      0
1  b  b      1
2  a  a      0
3  c  c      2

Thanks for the share. I am trying to tweak this to use for fields with high cardinality. Currently it gives memory error for fields with higher cardinality like 5000 or so. I read on stackoverflow and someone mentioned that the toarray() must be causing the issue.

what if you wanted to encode multiple columns simultaneously? Taking off from the above example, how could one encode the columns e and f in the following dataframe if you don't care whether a value appears in e or f, you just want to know if it appears at all?
df = pd.DataFrame(np.random.randn(25, 3), columns=['a', 'b', 'c'])

Make some random categorical columns

df['e'] = [random.choice(('Chicago', 'Boston', 'New York')) for i in range(df.shape[0])]
df['f'] = [random.choice(('Chicago', 'Boston', 'New York')) for i in range(df.shape[0])]

After you create new columns using get_dummies, consider you get e.Chicago and f.Chicago.
Now as you just want to know if Chicago appears at all irrespective of which column, just apply OR condition on both columns and create a new column and then drop the initial 2 columns. Not sure if there is a short cut for this.

From a machine learning perspective is there a preferred option between get_dummies and factorize. I have read that since factorize produces unequal distances between categorical values, that the vectorized output of get_dummies is preferred.

E.g.:

    red = 0
    blue = 1
    green = 2

    => green = 2* blue

which obviously makes no sense.

hi there, I found encode_onehot(df, cols) can only encode columns all of strings. When apply to df = pd.DataFrame({'category':[6,7,8,6,7,8], 'number':[1,2,3,4,5,6]}) the method vec.get_feature_names() will only return ['category'] and the encoding will fail.

add the method below will work well on both of the cases above:

def one_hot(df, cols):
    for each in cols:
        dummies = pd.get_dummies(df[each], prefix=each, drop_first=False)
        df = pd.concat([df, dummies], axis=1)
    return df

What is the best way to use sklearn's OneHotEncoder with pandas dataframe?