Pandas long to wide by subsets

Viewed 81

I have a dataframe of the following format

YR    WEEK  ACCOUNT   Tax   AMOUNT   SURCHARGE  Label
2021   1      1       10     5000     5          0
2021   1      2       20     10000    10         1
2021   2      1       20     2000      2         0
2021   2      2        5      10       3         1
2021   3      1       20     10000    10         1
2021   3      2       10     40000    50         1

I want to reshape it such that, it results in the following

ACCOUNT    YR    WEEK    Tax_1   Amount_1   Surcharge_1    Tax_2   Amount_2   Surcharge_2   Label
1         2021    2      10       5000         5             20     2000        2           0
2         2021    2      20       10000       10             5      10          3           1
1         2021    3      0         2000        2             20     10000       10          1
2         2021    3      5          10       3              10     40000        50          1

Each row is reshaped such that it is a function of the past 2 rows by account, year and week. And the Label is the current label of the WEEK I tried first to reshape it to wide

df.pivot('ACCOUNT', columns=['YR','WEEK'])

And was going to try a reshape after, but the pivotresulted in NaN values Any ideas on how I can achieve this in pandas Thx!

1 Answers

What you want is not a pivot, as pivot reshapes the data, while you here want to duplicate/shift it (you have more data cells after than before).

To take a simpler example, the goal is to go from something like this (6 values):

   id  group value
0   1      1     A
1   1      2     B
2   1      3     C
3   2      1     D
4   2      2     E
5   2      3     F

to this (8 values):

   id  group_1 value_1  group_2 value_2
0   1        1       A        2       B
1   1        2       B        3       C
2   2        1       D        2       E
3   2        2       E        3       F

How do we do it?

You can achieve what you want by doing a selection per group and concatenating the two outputs:

cols = ['YR', 'ACCOUNT']
g = df.sort_values(by=['YR', 'WEEK']).groupby(cols)

out = (
 pd.concat([g.apply(lambda d: d.drop(columns=cols+['WEEK']).iloc[:-1]).droplevel(-1).add_suffix('_1'),
            g.apply(lambda d: d.drop(columns=cols).iloc[1:]).droplevel(-1).add_suffix('_2')
            ], axis=1)
   .reset_index()
)

NB. the output is slightly different for the labels

output:

     YR  ACCOUNT  Tax_1  AMOUNT_1  SURCHARGE_1  Label_1  WEEK_2  Tax_2  AMOUNT_2  SURCHARGE_2  Label_2
0  2021        1     10      5000            5        0       2     20      2000            2        0
1  2021        1     20      2000            2        0       3     20     10000           10        1
2  2021        2     20     10000           10        1       2      5        10            3        1
3  2021        2      5        10            3        1       3     10     40000           50        1
slightly more condensed version

NB. This requires not to have NaNs

cols = ['YR', 'ACCOUNT']
g = df.sort_values(by=['YR', 'WEEK']).groupby(cols)
pd.concat([g.shift().add_suffix('_1'), df.add_suffix('_2')], axis=1).dropna(how='any', axis=0)
Related