Pandas - pivoting multiple columns into fewer columns with some level of detail kept

user_11569643 0
Viewed 306

Say I have the following code that generates a dataframe:

df = pd.DataFrame({"customer_code": ['1234','3411','9303'],
                   "main_purchases": [3,10,5],
                   "main_revenue": [103.5,401.5,99.0],
                   "secondary_purchases": [1,2,4],
                   "secondary_revenue": [43.1,77.5,104.6]
                  })

df.head()

img

There's the customer_code column that's the unique ID for each client.

And then there are 2 columns to indicate the purchases that took place and revenue generated from main branches by those clients.

And another 2 columns to indicate the purchases/revenue from secondary branches by those clients.

I want to get the data into a format like this, where a pivot is done where there's a new column to differentiate between main vs secondary, but the revenue numbers and purchase columns are not mixed up:

img2

The obvious solution is just to split this into 2 dataframes, and then simply do a concatenate, but I'm wondering whether there's a built-in way to do this in a line or two - this strikes me as the kind of thing someone might have thought to bake in a solution for.

pandas pivot python
3 Answers

With a little column renaming to get the "revenue" and "purchases" in the column names first using a regular expression and str.replace we can use pd.wide_to_long to convert these now stubnames from columns to rows:

# Reorder column names so stubnames are first
df.columns = [df.columns[0],
              *df.columns[1:].str.replace(r'(.*)_(.*)', r'\2_\1', regex=True)]

# Convert wide_to_long
df = (
    pd.wide_to_long(
        df,
        i='customer_code',
        stubnames=['purchases', 'revenue'],
        j='type',
        sep='_',
        suffix='.*'
    )
        .sort_index()  # Optional sort to match expected output
        .reset_index()  # retrieve customer_code from the index
)

df:

customer_code type purchases revenue
0 1234 main 3 103.5
1 1234 secondary 1 43.1
2 3411 main 10 401.5
3 3411 secondary 2 77.5
4 9303 main 5 99
5 9303 secondary 4 104.6

What does reordering the column headers do?

df.columns = [df.columns[0],
              *df.columns[1:].str.replace(r'(.*)_(.*)', r'\2_\1', regex=True)]

Produces:

Index(['customer_code', 'purchases_main', 'revenue_main',
       'purchases_secondary', 'revenue_secondary'],
      dtype='object')

The "type" column is now the suffix of the column header which allows wide_to_long to process the table as expected.

user_154978880

You can abstract the reshaping process with pivot_longer from pyjanitor; they are just a bunch of wrapper functions in Pandas:

#pip install pyjanitor
import pandas as pd
import janitor
df.pivot_longer(index = 'customer_code',
                names_to=('type', '.value'), 
                names_sep='_', 
                sort_by_appearance=True)
 
  customer_code       type  purchases  revenue
0          1234       main          3    103.5
1          1234  secondary          1     43.1
2          3411       main         10    401.5
3          3411  secondary          2     77.5
4          9303       main          5     99.0
5          9303  secondary          4    104.6

The .value in names_to signifies to the function that you want that part of the column to remain as a header; the other part goes under the type column. The split is determined in this case by names_sep (there is a names_pattern option, that allows regular expression split); if you do not care about the order of appearance, you can set sort_by_appearance as False.

user_71757130

You can also use melt() and concat() function to solve this problem.

import pandas as pd

df1 = df.melt(
          id_vars='customer_code',
          value_vars=['main_purchases', 'secondary_purchases'],
          var_name='type',
          value_name='purchases',
          ignore_index=True)

df2 = df.melt(
          id_vars='customer_code',
          value_vars=['main_revenue', 'secondary_revenue'],
          var_name='type',
          value_name='revenue',
          ignore_index=True)

Then we use concat() with the parameter axis=1 to join side by side and use sort_values(by='customer_code') to sort data by customer.

result= pd.concat([df1,df2['revenue']], 
               axis=1,
               ignore_index=False).sort_values(by='customer_code')

Using replace() with regex to align type names:

result.type.replace(r'_.*$','', regex=True, inplace=True)

The above code will output the below dataframe:

customer_code type purchases revenue
0 1234 main 3 103.5
3 1234 secondary 1 43.1
1 3411 main 10 401.5
4 3411 secondary 2 77.5
2 9303 main 5 99
5 9303 secondary 4 104.6
user_138581950
Related