Pandas create multiindex from array and OrderedDict-CodePudding

I have a pandas dataframe like so:

col1  col2  col3  col4  col5  col6
val1  val2  val3  val4  val5  val6
...

I have have an array having level1 column names and an OrderedDict having level0 column names and how many columns fall under them:

col_names = ['id', 'a', 'b', 'a', 'b', 'c']
col_layout_dict = OrderedDict([('ID', 1), ('A', 2), ('B', 2), ('C', 1)])

col_names and col_layout_dict are being used in other part of the code too, so I don't want to change them and also, since they are already available I would like to repurpose them to update the column names of my dataframe like so:

ID    A           B           C
id    a     b     a     b     c
val1  val2  val3  val4  val5  val6
...

How can I do this?

CodePudding user response：

Use list comprehension with range for flatten values of dict:

col_names = ['id', 'a', 'b', 'a', 'b', 'c']
col_layout_dict = OrderedDict([('ID', 1), ('A', 2), ('B', 2), ('C', 1)])

c = [k for k, v in col_layout_dict.items() for x in range(v)]

col_names = ['id', 'a', 'b', 'a', 'b', 'c']

mux = pd.MultiIndex.from_arrays([c, col_names])
print (mux)
MultiIndex([('ID', 'id'),
            ( 'A',  'a'),
            ( 'A',  'b'),
            ( 'B',  'a'),
            ( 'B',  'b'),
            ( 'C',  'c')],
           )

 df.columns = mux