Match column to another column containing array

Viewed 33

I have very junior question in python - i have a dataframe with a column containing some IDs and separate dataframe that contains 2 columns, out of which 1 is an array:

df1 = pd.DataFrame({"some_id": [1, 2, 3, 4, 5]})
df2 = pd.DataFrame([["A", [1, 2]], ["B", [3, 4]], ["C", [5]]], columns=['letter', 'some_ids'])

I want to add do df1 new column "letter' that for a given "some_id" will look up df2, check if this id is in df2['some_ids'] and return df2['letter']

I tried this:

df1['letter'] = df2[df1[some_id].isin(df2['some_ids')].letter

and get NaNs - any suggestion where I make mistake?

1 Answers

Create dictionary with flatten nested lists in dict comprehension and then use Series.map:

d = {x: a for a,b in zip(df2['letter'], df2['some_ids']) for x in b}
df1['letter'] = df1['some_id'].map(d)

Or mapping by Series created by DataFrame.explode with DataFrame.set_index:

df1['letter'] = df1['some_id'].map(df2.explode('some_ids').set_index('some_ids')['letter'])

Or use left join with rename column:

df1 = df1.merge(df2.explode('some_ids').rename(columns={'some_ids':'some_id'}), how='left')

print (df1)
  some_id letter
0       1      A
1       2      A
2       3      B
3       4      B
4       5      C
Related