Search code examples
pythonpandasdataframesklearn-pandas

encode list in a cell in a dataframe with sklearn


I have a data frame like this:

df = pd.DataFrame([{'A': 1, 'B': 'p'}, {'A': 1, 'B': 'q'},  {'A': 2, 'B': 'o'},  {'A': 3, 'B': 'p'}])
df
   A  B
0  1  p
1  1  q
2  2  o
3  3  p

I could encode and decode it correctly with with code.

le = LabelEncoder()
df_encoded = pd.DataFrame(columns=df.columns)
df_decoded = pd.DataFrame(columns=df.columns)

for col in df.columns:
    df_encoded[col] = le.fit_transform(df[col])

df_encoded
   A  B
0  0  1
1  0  2
2  1  0
3  2  1

for col in df.columns:
    le = le.fit(df[col])
    df_decoded[col] = le.inverse_transform(df_encoded[col])

df_decoded

   A  B
0  1  p
1  1  q
2  2  o
3  3  p  

Now if I have a data frame like this, how can I encode and decode it?

dj = pd.DataFrame([{'A': [1,2], 'B': 'p'}, {'A': 1, 'B': ['p','q']},  {'A': 2, 'B': 'o'},  {'A': 3, 'B': 'p'}])

I want to have a code for each cell of ['p','q'] instead of a code for ['p','q'].


Solution

  • One way to do it is to break down the cells that include lists into separate rows, then apply LabelEncoder, and then combine these rows back into lists:

    df_encoded = pd.DataFrame()
    df_decoded = pd.DataFrame()
    
    def t1(z):
        zz = pd.DataFrame([np.array(x).reshape(-1) for x in z.values.tolist()])
        dt = zz.dtypes[0]
        return (zz
            .stack()
            .reset_index(level=1, drop=True)
            .to_frame(col)
            .astype(dt))
    
    def t2(z):
        return z.groupby(level=0).apply(lambda x: np.squeeze(x.values.tolist()))
    
    for col in dj.columns:
        d = t1(dj[col])
        d['x'] = le.fit_transform(d[col])
        df_encoded[col] = t2(d['x'])
    
    print(df_encoded)
    
    for col in dj.columns:
        d = t1(dj[col])
        m = le.fit(d[col])
    
        d = t1(df_encoded[col])
        d['x'] = m.inverse_transform(d[col])
        df_decoded[col] = t2(d['x'])
    
    print(df_decoded)
    

    Output:

            A       B
    0  [0, 1]       1
    1       0  [1, 2]
    2       1       0
    3       2       1
    
            A       B
    0  [1, 2]       p
    1       1  [p, q]
    2       2       o
    3       3       p