Pandas: creating a new column conditional on substring searches of one column and inverse of another column

user_5179643 0
Viewed 1992

I'd like to create a new column in a Pandas data frame based on a substring search of one column and an inverse of another column. Here is some data:

import pandas as pd
import numpy as np

df = pd.DataFrame({'Manufacturer':['ABC-001', 'ABC-002', 'ABC-003', 'ABC-004', 'DEF-123', 'DEF-124', 'DEF-125', 'ABC-987', 'ABC-986', 'ABC-985'],
                   'Color':['04-Red', 'vs Red - 07', 'Red', 'Red--321', np.nan, np.nan, np.nan, 'Blue', 'Black', 'Orange'],
                  })


    Manufacturer    Color
0   ABC-001         04-Red
1   ABC-002         vs Red - 07
2   ABC-003         Red
3   ABC-004         Red--321
4   DEF-123         NaN
5   DEF-124         NaN
6   DEF-125         NaN
7   ABC-987         Blue
8   ABC-986         Black
9   ABC-985         Orange

I would like to be able to create a new column named Country based on the following logic:

a) if the Manufacturer column contains the substring 'ABC' and the Color column contains the substring 'Red', then write 'United States' to the Country column

b) if the Manufacturer column contains the substring 'DEF', then write 'Canada to the Country column

c) if the Manufacturer column contains the substring 'ABC' and the Color column does NOT contain the substring 'Red', then write 'England' to the Country column.

My attempt is as follows:

df['Country'] = np.where((df['Manufacturer'].str.contains('ABC')) & (df['Color'].str.contains('Red', na=False)), 'United States',  # the 'a' case
                np.where(df['Manufacturer'].str.contains('DEF', na=False), 'Canada',  # the 'b' case
                np.where((df['Manufacturer'].str.contains('ABC')) & (df[~df['Color'].str.contains('Red', na=False)]), 'England',  # the 'c' case
                         'ERROR')))

But, this gets the following error:

TypeError: Cannot perform 'rand_' with a dtyped [float64] array and scalar of type [bool]

The error message suggests that it might be a matter of operator precedence, as mentioned in:

pandas comparison raises TypeError: cannot compare a dtyped [float64] array with a scalar of type [bool]

Python error: TypeError: cannot compare a dtyped [float64] array with a scalar of type [bool]

I believe I'm using parentheses properly here (although maybe I'm not).

Does anyone see the cause of this error? (or know of a more elegant want to accomplish this?)

Thanks in advance!

pandas numpy python
3 Answers

You don't want to index into df here, so just do this:

Just change: (df[~df['Color'].str.contains('Red', na=False)])

to: ~df['Color'].str.contains('Red', na=False)

and it should work.

Also, if you want to break this up for readability and to eliminate some repetition, I would suggest something like this:

# define the parameters that define the Country variable in another table
df_countries = pd.DataFrame(
    {'letters': ['ABC', 'DEF', 'ABC'],
     'is_red': [True, False, False],
     'Country': ['United States', 'Canada', 'England']})

# add those identifying parameters to your current table as temporary columns
df['letters'] = df.Manufacturer.str.replace('-.*', '')
df['is_red'] = df.Color.str.contains('Red', na=False)

# merge the tables together and drop the temporary key columns
df = df.merge(df_countries, how='left', on=['letters', 'is_red'])
df = df.drop(columns=['letters', 'is_red'])

Or more concise:

in_col = lambda col, string: df[col].str.contains(string, na=False)

conds = {'United States': in_col('Manufacturer', 'ABC') & in_col('Color', 'Red'),
         'Canada': in_col('Manufacturer', 'DEF'),
         'England': in_col('Manufacturer', 'ABC') & ~in_col('Color', 'Red')}

df['Country'] = np.select(condlist=conds.values(), choicelist=conds.keys())
user_25662830

Another way out is use of np.select(list of conditions, list of choices)

conditions=[(df['Manufacturer'].str.contains('ABC')) & (df['Color'].str.contains('Red')),df['Manufacturer'].str.contains('DEF', na=False),(df['Manufacturer'].str.contains('ABC')) & (~df['Color'].str.contains('Red', na=False))]
choices=['United States','Canada','England']
df['Country']=np.select(conditions,choices)



  Manufacturer        Color        Country
0      ABC-001       04-Red  United States
1      ABC-002  vs Red - 07  United States
2      ABC-003          Red  United States
3      ABC-004     Red--321  United States
4      DEF-123          NaN         Canada
5      DEF-124          NaN         Canada
6      DEF-125          NaN         Canada
7      ABC-987         Blue        England
8      ABC-986        Black        England
9      ABC-985       Orange        England
user_89869750

This is an easy and straightforward way to do it:

country = []

for index, row in df.iterrows():
    
    if 'DEF' in row['Manufacturer']:
        country.append('Canada')
    elif 'ABC' in row['Manufacturer']:
        if 'Red' in row['Color']:
            country.append('United States')
        else:
            country.append('England')
    else:
        country.append('')
        
df['Country'] = country

Of course there will be more efficient ways to go about this w/o looping through the entire dataframe, but in almost all cases this should be sufficient.

user_121447870
Related