Check dataframe value to dictionary, append key to new column-CodePudding

I am trying to check a row value for a match in a dictionary and append the key to a new column.

Example:

group_ID = {
    'Group A':[4738, 4812],
    'Group B':[5888, 6551],
    'Group C':[4487, 7888]
}

user_data = [['Alex',4812],['Bob',4487],['Clarke',5888]]
sample_df = pd.DataFrame(user_data,columns=['Name','User ID'])
print(sample_df)

     Name  User ID
0    Alex     4812
1     Bob     4487
2  Clarke     5888

Using this example, if 'User ID' in sample_df has a matching value in dictionary 'group_ID' then I would like to add a third column reflecting the key name like below:

     Name  User ID Group ID
0    Alex     4812  Group A
1     Bob     4487  Group C
2  Clarke     5888  Group B

Thanks in advance for the help!

CodePudding user response：

Does this do the job:

matches = []
# iterate over each user ID
for user_id in sample_df['User ID']:
    # go through the items in the dictionary and append to `matches` list
    matches.extend([key for key, values in group_ID.items() if user_id in values])

if matches:
    # if matches not empty, then add a new column
    sample_df['group_ID'] = matches

CodePudding user response：

Try:

tmp = {i: k for k, v in group_ID.items() for i in v}
sample_df["Group ID"] = sample_df["User ID"].map(tmp)
print(sample_df)

Prints:

     Name  User ID Group ID
0    Alex     4812  Group A
1     Bob     4487  Group C
2  Clarke     5888  Group B