Search code examples
pythonpandasdataframepivot-tablepandas-melt

Reshape DataFrame Pandas - some variables to long others to wide


I need to reshape a dataframe so that some of the variables (Diag1, Diag2, Diag3) change to long wile others (Period) change to wide. Basically they need to swap places.

I have recreated the original dataframe in the example below. I've tried using pivot and melt separately to no avail as demonstrated in the example.

df = pd.DataFrame({
    'ID':[1,1,2,2,3,3,4,4,5,5,6,6,7,7,8,8,9,9,10,10,11,11,],
    'Period':['0 Month','3 Month','0 Month','3 Month','0 Month',
            '3 Month','0 Month','3 Month','0 Month','3 Month','0 Month',
            '3 Month','0 Month','3 Month','0 Month','3 Month','0 Month',
            '3 Month','0 Month','3 Month','0 Month','3 Month',],
    'Diag1':[0,1,1,0,0,0,0,1,0,0,1,0,0,0,0,0,0,0,1,0,0,0,],
    'Diag2':[0,0,1,0,0,0,0,1,0,0,1,0,0,0,0,0,0,0,1,0,0,0,],
    'Diag3':[0,0,1,0,0,0,1,1,0,0,1,0,0,0,1,0,0,0,1,0,1,0,]
    })

dfp = df.pivot(index=["ID",], columns='Period',).reset_index()
print(dfp)

dfm = df.melt(id_vars=["ID",],value_vars=['Period'])
print(dfm)

The desired outcome is:

ID  Diagnosis   0_Month 3_Month
1   Diag1       0       1
1   Diag2       0       0
1   Diag3       0       0
2   Diag1       0       1
2   Diag2       1       0
2   Diag3       1       0
3   Diag1   

3   Diag2   

3   Diag3       etc...

I suspect I need some combination of the 2 but am struggling to find any examples. My brain is beginning to melt as a consequence...


Solution

  • You could melt; then pivot:

    out = (df.melt(id_vars=['ID', 'Period'], var_name='Diagnosis')
           .pivot(['ID','Diagnosis'], 'Period', 'value')
           .reset_index().rename_axis(columns=[None]))
    

    Output:

        ID Diagnosis  0 Month  3 Month
    0    1     Diag1        0        1
    1    1     Diag2        0        0
    2    1     Diag3        0        0
    3    2     Diag1        1        0
    4    2     Diag2        1        0
    ..  ..       ...      ...      ...
    28  10     Diag2        1        0
    29  10     Diag3        1        0
    30  11     Diag1        0        0
    31  11     Diag2        0        0
    32  11     Diag3        1        0
    
    [33 rows x 4 columns]