Pandas/Python create columns based on multiple conditions with a dictionary

Question

I am looking for a shorter and more elegant way of writing the condition below, let's say I have 30 different combinations. Is it possible to map some kind of dictionary with multiple conditions or something else?

def set_name(df):
    if df['PRODUCT'] == 'A' and df['TYPE'] == 'C2':
        return 'AAC2'
    elif df['PRODUCT'] == 'A' and df['TYPE'] == 'C3':
        return 'AAC3'
    elif df['PRODUCT'] == 'B' and df['TYPE'] == 'C2':
        return 'BBC"'
    
df['RETURN'] = df.apply(set_name, axis=1)

Another way is the below which still pretty long if I have many combinations.

conditions = [
    (df['PRODUCT'] == 'A') & (df['TYPE'] == 'C2'),
    (df['PRODUCT'] == 'A') & (df['TYPE'] == 'C3')
]

values = ['AAC2', 'AAC3']

df['RETURN'] = np.select(conditions, values)

score 3 · Answer 1 · answered Jan 11 '21 at 08:36

Create helper DataFrame and use left join by DataFrame.merge:

products = ['A','A']
types = ['C2','C3']
values = ['AAC2', 'AAC3']

df1 = pd.DataFrame({'PRODUCT':products,
                    'TYPE':types,
                    'RETURN': values})


df = df.merge(df1, on=['PRODUCT','TYPE'], how='left')

score 2 · Answer 2 · answered Jan 11 '21 at 19:00

You could use a two level dictionary or a dictionary with (product,type) tuples as key:

Two level dictionary ...

mapping = {
            'A': {'C2':'AAC2', 'C3':'AAC3'},
            'B': {'C2':'BBC'}
          } 
            
def set_name(df):
    product, ptype = df['PRODUCT'], df['TYPE']
    if product in mapping and ptype in mapping[product]:
        return mapping[product][ptype]
    else:
        # handle no-mapping situations
                                   

    
df['RETURN'] = df.apply(set_name, axis=1)

Or, using tuple keys ...

mapping = {
            ('A','C2'): 'AAC2',
            ('A','C3'): 'AAC3',
            ('B','C2'): 'BBC'
          } 
            
def set_name(df):
    product, ptype = df['PRODUCT'], df['TYPE']
    if (product,ptype) in mapping:
        return mapping[product,ptype]
    else:
        # handle no-mapping situations
                                
    
df['RETURN'] = df.apply(set_name, axis=1)

Pandas/Python create columns based on multiple conditions with a dictionary

2 Answers2