How to More Efficiently Pivot Semi Colon Separated Columns into a 0/1/2 Indicator Matrix?

user_16052735 0
Viewed 68

I am starting with a pandas Dataframe that is generated by the following code:

import pandas as pd

data = {'basket_1':['apple;banana;orange', 'apple;banana;mango', 'mango;orange;grapefruit'],
        'basket_2':['pineapple;strawberry;peach', 'peach;lemon;guava', 'strawberry;peach;guava']}
  
# Create DataFrame
df = pd.DataFrame(data)


                  basket_1                    basket_2
0      apple;banana;orange  pineapple;strawberry;peach
1       apple;banana;mango           peach;lemon;guava
2  mango;orange;grapefruit      strawberry;peach;guava

I also start with a numpy array of all of the fruits present in either basket column (as well as some additional fruits), generated as follows for the purpose of this example:

import numpy as np

fruits = np.array(['apple', 'banana', 'orange', 'mango', 'grapefruit', 'kiwi', 'pineapple', 
                   'strawberry', 'peach', 'lemon', 'guava', 'lime'])

From this initial Dataframe and array, I am looking to generate the resulting Dataframe in the most efficient way possible:

Final Dataframe

                  basket_1                    basket_2  apple  banana  orange  mango  grapefruit  kiwi  pineapple  strawberry  peach  lemon  guava  lime
0      apple;banana;orange  pineapple;strawberry;peach      1       1       1      0           0     0          2           2      2      0      0     0
1       apple;banana;mango           peach;lemon;guava      1       1       0      1           0     0          0           0      2      2      2     0
2  mango;orange;grapefruit      strawberry;peach;guava      0       0       1      1           1     0          0           2      2      0      2     0

The final result has added a column for each of the elements present in the fruits array along with a '1' in a given row if the given fruit is present in the 'basket_1' column, a '2' in a given row if the given fruit is present in the 'basket_2' column, and a '0' otherwise.

At the moment, I am using the following code to transform the initial Dataframe into the desired format:

def whichBasket(b1, b2, fruit):
    if fruit in b1:
        val = 1
    elif fruit in b2:
        val = 2
    else:
        val = 0
    return val

for f in fruits:
    df[f] = df.apply(lambda x: whichBasket(x.basket_1, x.basket_2, f), 1)

This solution calls an apply function iterating through each row of the Dataframe nested inside a for-loop iterating through each fruit, which works fine for a small example such as this one. However, I am attempting to scale this up to a Dataframe with over 1000 fruits and over 80000 rows, and this solution is far too slow to complete this job in a reasonable amount of time.

Any ideas for ways to improve this code's performance in terms of shortening running time? Thanks in advance for your help.

pandas pivot lambda python performance
2 Answers

Here is one option,

  1. create dataframe by appending all the columns from fruits list to current df

  2. User str.get_dummies to generate dummy columns and assign to original df

  3. For basket 2, add 1 to ensure str.get_dummies returns 2 as value instead of 1

  4. Replace nan in missing cells with 0

     df = pd.concat([df, pd.DataFrame(columns = fruits)], axis = 1)

 df = df.assign(**df['basket_1'].str.get_dummies(';'))

 df = df.assign(**df['basket_2'].str.get_dummies(';') * 2)

 df = df.fillna(0)

Edit: addressing @Matt's question on a fruit being present in either basket but not at the same time. Slightly different approach, concat df with get_dummies results from baskets 1 & 2. Use reindex to include rest of the columns from fruits list.

data = {'basket_1':['apple;banana;orange', 'apple;banana;mango', 'mango;orange;grapefruit','mango;peach;grapefruit'],
        'basket_2':['pineapple;strawberry;peach', 'peach;lemon;guava', 'strawberry;peach;guava', 'apple;lemon;guava']}
  
df = pd.DataFrame(data)

fruits = np.array(['apple', 'banana', 'orange', 'mango', 'grapefruit', 'kiwi', 'pineapple', 
                   'strawberry', 'peach', 'lemon', 'guava', 'lime'])
req_cols = df.columns.tolist() + fruits.tolist()

df = pd.concat([df, df['basket_1'].str.get_dummies(';'), df['basket_2'].str.get_dummies(';')*2], axis = 1).groupby(level=0, axis=1).sum()
df = df.reindex(req_cols, axis = 1, fill_value = 0)
user_62873080

You can perform:

  1. Get basket_1 entries with str.split() and then create columns with pd.Series()
  2. Similarly, for basket_2
  3. Concat df with the 2 sets of newly created columns from the 2 baskets
  4. Reindex new columns with list of fruits and fillna() with 0 for non-defined entries.
df_b1 = df['basket_1'].str.split(';').apply(lambda x: pd.Series([1] * len(x), index=x))
df_b2 = df['basket_2'].str.split(';').apply(lambda x: pd.Series([2] * len(x), index=x))
df1 = pd.concat([df, df_b1, df_b2],  axis=1)
df1 = df1.set_index(['basket_1', 'basket_2']).reindex(columns=fruits).fillna(0, downcast='infer').reset_index()

Result:

print(df1)


                  basket_1                    basket_2  apple  banana  orange  mango  grapefruit  kiwi  pineapple  strawberry  peach  lemon  guava  lime
0      apple;banana;orange  pineapple;strawberry;peach      1       1       1      0           0     0          2           2      2      0      0     0
1       apple;banana;mango           peach;lemon;guava      1       1       0      1           0     0          0           0      2      2      2     0
2  mango;orange;grapefruit      strawberry;peach;guava      0       0       1      1           1     0          0           2      2      0      2     0
user_150706970
Related