Pandas reshape extracting multiple values from colname

Viewed 40

I have a wide dataframe I want to be able to reshape.

I have some columns that I wanna preserve. I have been exploring melt and wide_to_long but I'm not sure that's what I need.

Imagine I have some columns named: 'id', 'classroom', 'city' And other columns called: 'alumn_x_subject_y_mark', 'alumn_x_subject_y_name', 'alumn_x_subject_y_teacher'

And x and y are the product of [range(20), range(10)].

I would like to end with a df that has columns: id, classroom, city, alumn, subject, mark, name, teacher

With all the original 20*10 columns converted to rows.

An empty dataframe with that structure can be generated this way:

import pandas as pd
import itertools
vals = list(itertools.product(*[range(20), range(10)]))
pd.DataFrame(columns=['id', 'classroom', 'city']+ \
             ['alumn_{0}_subject_{1}_mark'.format(x, y) for x, y in vals] + \
             ['alumn_{0}_subject_{1}_name'.format(x, y) for x, y in vals] + \
             ['alumn_{0}_subject_{1}_teacher'.format(x, y) for x, y in vals]
             , dtype=object)

I'm not building this dataframe but receiving it from a file, that's why it has so many columns and I cannot change that.

1 Answers

If you had only 2 parameters to extract, wide_to_long would work.

Here you have 3, thus you can perform a manual reshaping with a MultiIndex:

regex = r'alumn_(\d+)_subject_(\d+)_(.*)'

out = (df
 .set_index(['id', 'classroom', 'city'])
 .pipe(lambda d: d.set_axis(pd.MultiIndex
                              .from_frame(d.columns.str.extract(regex),
                                          names=['alumn', 'subject', None]
                                          ),
                            axis=1))
 .stack(['alumn', 'subject'])
 .reset_index()
)

output:

Empty DataFrame
Columns: [id, classroom, city, alumn, subject, mark, name, teacher]
Index: []

output with a single row (after df.loc[0] = range(df.shape[1])):

     id  classroom  city alumn subject mark name teacher
0     0          1     2     0       0    3  203     403
1     0          1     2     0       1    4  204     404
2     0          1     2     0       2    5  205     405
3     0          1     2     0       3    6  206     406
4     0          1     2     0       4    7  207     407
..   ..        ...   ...   ...     ...  ...  ...     ...
195   0          1     2     9       5   98  298     498
196   0          1     2     9       6   99  299     499
197   0          1     2     9       7  100  300     500
198   0          1     2     9       8  101  301     501
199   0          1     2     9       9  102  302     502

[200 rows x 8 columns]
Related