Home > Software engineering >  Pandas: Pivot multi-index, with one 'shared' column
Pandas: Pivot multi-index, with one 'shared' column

Time:11-20

I have a pandas dataframe that can be represented like:

test_dict = {('a', 1) : {'shared':0,'x':1, 'y':2, 'z':3}, 
            ('a', 2) : {'shared':1,'x':2, 'y':4, 'z':6},
            ('b', 1) : {'shared':0,'x':10, 'y':20, 'z':30},
            ('b', 2) : {'shared':1,'x':100, 'y':200, 'z':300}}
example = pd.DataFrame.from_dict(test_dict).T

enter image description here

I am trying to figure out a way to turn this into a dataframe that looks like this dictionary representation:

res_dict = {1 : {'shared':0,'a':{'x':1, 'y':2, 'z':3}, 'b':{'x':10, 'y':20, 'z':30}},
            2 : {'shared':1,'a':{'x':2, 'y':4, 'z':6},'b':{'x':100, 'y':200, 'z':300}}}
        

Any suggestions appreciated! Thanks

CodePudding user response:

A possible solution, which uses only dataframe manipulations and then converts to dictionary:

xyz = ['x', 'y', 'z']
out = (example.assign(xyz=example[xyz].apply(list, axis=1)).reset_index()
       .pivot(index='level_0', columns=['level_1', 'shared'], values='xyz')
       .applymap(lambda x: dict(zip(xyz, x))))

out.columns = out.columns.rename(None, level=0)
out.index = out.index.rename(None)

(pd.concat([out.droplevel(1, axis=1), 
           out.columns.to_frame().reset_index(drop=True).iloc[:,1]
           .to_frame().T.set_axis(out.columns.get_level_values(0), axis=1)])
 .iloc[np.arange(-1, len(out))].to_dict())

Output:

{
    1: {
        'shared': 0,
        'a': {'x': 1, 'y': 2, 'z': 3},
        'b': {'x': 10, 'y': 20, 'z': 30}
    },
    2: {
        'shared': 1,
        'a': {'x': 2, 'y': 4, 'z': 6},
        'b': {'x': 100, 'y': 200, 'z': 300}
    }
}
  • Related