Pandas dataframe look for key in dict, write key in new column according to value

user_13817966 0
Viewed 66

I have a dataframe:

df = pd.DataFrame([
                   {'ID': 1,'A': [{'name': 'lifestyle'}, {'name': 'economy'}, 
                          {'name': 'politics'}, {'name': 'climate & environment'}]}, 
                   {'ID': 2,'A': [{'name': 'sport'}]}, 
                   {'ID': 3,'A': [{'name': 'climate & environment'}]},
                   {'ID': 4,'A': [{'name': 'sport'}]},
                   {'ID': 5,'A': [{'name': 'politics'}, {'name': 'world'}]},
                   {'ID': 6,'A': [{'name': 'economy'}, {'name': 'politics'}]}
                  ])

Every value in col A belongs to a category. Those categories are in a hardcoded dict (categories.txt):

dict= {'lifestyle':'cat1',
'economy':'cat1',
'politics':'cat2',
'climate & environment':'cat2',
'sport':'cat3',
'world':'cat4',
'news':'cat3'}

My goal is to look for every key and write this key(s) in a new col named after the value (cat1, cat2,...)

Here is what I got so far:

df['A'] = [','.join(map(str, l)) for l in df['A']]

# read in the dict
d = {}
with open("categories.txt", "r") as file:
    for line in file:
        key, value = line.strip().split(":")
        d[key] = value


di = {k: oldk for oldk, oldv in d.items() for k in oldv.split(',')}


for k, v in d.items():
    if v == 'cat1':
        df.loc[df['A'].str.contains(k), 'cat1'] = k
    elif v == 'cat2':
        df.loc[df['A'].str.contains(k), 'cat2'] = k
    elif v == 'cat3':
        df.loc[df['A'].str.contains(k), 'cat3'] = k 
    else:
        df.loc[df['A'].str.contains(k), 'cat4'] = k

Right now, if there is more than one key per row this will be written as the next category which is not the right way. How do I get every key (one or multiple keys per cell) in the right column named after the values? Something like this:

df = pd.DataFrame({'ID':[1,2,3,4,5,6],
                   'cat1':['lifestyle, economy','nan','nan','nan','nan','economy, politics'],
                   'cat2':['politics, climate & environment','nan','climate & environment','nan','politics','nan'],
                   'cat3':['nan','sport','nan','sport','nan','nan'],
                   'cat4':['nan','nan','nan','nan','world','nan']})

Thanks in advance

pandas dataframe key-value dictionary python
1 Answers

Use explode and apply to extract values from your first dataframe then map with your dict before pivot:

Update

Merging the ID col back with the categories does not work because the amount of entries is different. The ID col is crucial

out = df['A'].explode().apply(pd.Series).reset_index()
out['category'] = out['name'].map(d)
out = out.pivot_table(index='index', columns='category',
                      values='name', aggfunc=', '.join) \
         .rename_axis(index=None, columns=None)
out = df[['ID']].join(out)

Output result:

>>> out
   ID                cat1                             cat2   cat3   cat4
0   1  lifestyle, economy  politics, climate & environment    NaN    NaN
1   2                 NaN                              NaN  sport    NaN
2   3                 NaN            climate & environment    NaN    NaN
3   4                 NaN                              NaN  sport    NaN
4   5                 NaN                         politics    NaN  world
5   6             economy                         politics    NaN    NaN
user_152399510
Related