pythonpandasdataframedictionarymapping

Pandas Map Dictionary with Multiple Values to Dataframe


I have the following dataframe:

data = [['BALANCED', 'Corp EUR'], ['YIELD', 'Corp USD'], ['GROWTH', 'HG CHF']]
df = pd.DataFrame(data, columns=['STRATEGY', 'ASSET_CLASS']
df

    STRATEGY    ASSET_CLASS
0   BALANCED    Corp EUR
1   YIELD       Corp USD
2   GROWTH      HG CHF

My objective is to assign a specific value from a dictionary with multiple values depending on the value in column "STRATEGY". The dictionary containing the different asset classes looks as follows:

my_dict = {'HG CHF': [5, 2, 3, 4],
           'Corp EUR': [4, 6, 8, 7],
           'Corp USD': [9, 7.5, 5.3, 6],
          }
my_dict

{'HG CHF': [5, 2, 3, 4],
 'Corp EUR': [4, 6, 8, 7],
 'Corp USD': [9, 7.5, 5.3, 6]}

In the dictionary, each value depicts a specific strategy. In this case, YIELD include all second values, BALANCED all third values and GROWTH all fourth values, respectively. Essentially, what I'd like to obtain is the following output:

   STRATEGY ASSET_CLASS  TARGET
0  BALANCED    Corp EUR     8.0
1     YIELD    Corp USD     7.5
2    GROWTH      HG CHF     4.0

How can I correctly tell Python which value in the dictionary to access? Below command maps all values:

df['TARGET'] = df['ASSET_CLASS'].map(my_dict)

   STRATEGY ASSET_CLASS            TARGET
0  BALANCED    Corp EUR      [4, 6, 8, 7]
1     YIELD    Corp USD  [9, 7.5, 5.3, 6]
2    GROWTH      HG CHF      [5, 2, 3, 4]

Any suggestion is much appreciated!! Thanks a lot in advance!!


Solution

  • Use a second dictionary and a list comprehension with zip:

    # which item should be picked?
    nth = {'YIELD': 1, 'BALANCED': 2, 'GROWTH': 3}
    
    df['TARGET'] = [
        my_dict[a][nth[s]] for s, a in zip(df['STRATEGY'], df['ASSET_CLASS'])
    ]
    

    Output:

       STRATEGY ASSET_CLASS  TARGET
    0  BALANCED    Corp EUR     8.0
    1     YIELD    Corp USD     7.5
    2    GROWTH      HG CHF     4.0
    

    If there is a chance that the lists don't have enough items, you can add a safety check:

    # let's pick the 5th item for "GROWTH"
    nth = {'YIELD': 1, 'BALANCED': 2, 'GROWTH': 5}
    
    df['TARGET'] = [
        d[n] if (n := nth[s]) < len(d := my_dict[a]) else None
        for s, a in zip(df['STRATEGY'], df['ASSET_CLASS'])
    ]
    

    Example:

       STRATEGY ASSET_CLASS  TARGET
    0  BALANCED    Corp EUR     8.0
    1     YIELD    Corp USD     7.5
    2    GROWTH      HG CHF     NaN