Search code examples
pythonpandashierarchical-data

Replace columns with multi level columns based on lookup on Dataframe


How can I replace this single column heading:

    foo bar
0   0   0
1   0   0

To get these multi level columns:

    A   B
    a   b
0   0   0
1   0   0

Based on this dataframe mapping

    col1    col2    col3
0   foo     a       A
1   bar     b       B
2   baz     c       C

I am trying a list comprehension trying to create a new multi level column index but doesn't seem to be working...I have a feeling there is a more pythonic way to achieve this nonetheless

df1 = pd.DataFrame({'foo':[0,0],
                    'bar':[0,0]})
df2 = pd.DataFrame({'col1':['foo','bar','baz'],
                    'col2':['A','B','C'],
                    'col3':['a','b','c']})

df1.columns = [(df2.loc[df2['col1']==i,'col2'], df2.loc[df2['col1']==i,'col3']) for i in df1.columns]

Solution

  • You can transform df2 to a Series of tuples and map it to the columns:

    df1.columns = df1.columns.map(df2.set_index('col1').apply(tuple, axis=1))
    

    output:

       A  B
       a  b
    0  0  0
    1  0  0