Create new df column from dictionary value on matching key

February 24, 2022

Given the DataFrame df and dictionary test_dict:

import pandas as pd
 
# initialize data of lists.
data = {'ref':['Gfg', 'Gfg', 'Gfg', 'Gfg', 'isR', 'Qst', 'jPu'],
    'position':[1, 3, 1, 1, 2, 1, np.nan]}
 
# Create DataFrame
df = pd.DataFrame(data)

   ref  position
0  Gfg       1.0
1  Gfg       3.0
2  Gfg       1.0
3  Gfg       1.0
4  isR       2.0
5  Qst       1.0
6  jPu       NaN

test_dict = {"Gfg" : [5, 3, 6], "isR" : [8, 4], "Qst" : [10, 11]}

I am trying to create a new column in df by matching df.ref item against dictionary key and then returns the value from the paired list in test_dict based on value in df.position

This would produce the following:

   ref  position  new_col
0  Gfg       1.0      5
1  Gfg       3.0      6
2  Gfg       1.0      5
3  Gfg       1.0      5
4  isR       2.0      4
5  Qst       1.0      10
6  jPu       NaN      NaN

>Solution :

Here’s one approach:

map test_dict to "ref" column and use a list comprehension to iterate over the lists in each row and index them using "position":

df['new_col'] = [lst if isinstance(lst, float) else lst[int(i)-1] for i, lst in zip(df['position'], df['ref'].map(test_dict))]

Output:

   ref  position  new_col
0  Gfg       1.0      5.0
1  Gfg       3.0      6.0
2  Gfg       1.0      5.0
3  Gfg       1.0      5.0
4  isR       2.0      4.0
5  Qst       1.0     10.0
6  jPu       NaN      NaN