Search code examples
pythonstringpandasdictionarystring-matching

Pandas - check if dataframe columns contain key:value pairs from a dictionary


This question is related to another question I had posted. Pandas - check if a string column in one dataframe contains a pair of strings from another dataframe

My goal is to check if two different columns of a dataframe contain a pair of string values and if the condition is met, then extract one of the values.

I have two dataframes like this:

df1 = pd.DataFrame({'consumption':['squirrelate apple', 'monkey likesapple', 
                                  'monkey banana gets', 'badger/getsbanana', 'giraffe eats grass', 'badger apple.loves', 'elephant is huge', 'elephant/eats/', 'squirrel.digsingrass'], 
                    'name': ['apple', 'appleisred', 'banana is tropical', 'banana is soft', 'lemon is sour', 'washington apples', 'kiwi', 'bananas', 'apples']})

df2 = pd.DataFrame({'food':['apple', 'apple', 'banana', 'banana'], 'creature':['squirrel', 'badger', 'monkey', 'elephant']})

In [187]:df1
Out[187]: 
            consumption                name
0     squirrelate apple               apple
1     monkey likesapple          appleisred
2    monkey banana gets  banana is tropical
3     badger/getsbanana      banana is soft
4    giraffe eats grass       lemon is sour
5    badger apple.loves   washington apples
6      elephant is huge                kiwi
7        elephant/eats/             bananas
8  squirrel.digsingrass              apples

In[188]: df2
Out[188]: 
   creature    food
0  squirrel   apple
1    badger   apple
2    monkey  banana
3  elephant  banana

What I want to do is test if 'apple' occurs in df1['name'] and 'squirrel' occurs in df1['consumption'] and if both conditions are met then extract 'squirrel' from df1['consumption'] into a new column df['creature']. The result should look like:

Out[189]: 
            consumption  creature                name
0     squirrelate apple  squirrel               apple
1     monkey likesapple       NaN          appleisred
2    monkey banana gets    monkey  banana is tropical
3     badger/getsbanana       NaN      banana is soft
4    giraffe eats grass       NaN       lemon is sour
5    badger apple.loves    badger   washington apples
6      elephant is huge       NaN                kiwi
7        elephant/eats/  elephant             bananas
8  squirrel.digsingrass       NaN              apples

If there was no paired value constraint, I could have done something simple like :

np.where((df1['consumption'].str.contains(<creature_string>, case = False)) & (df1['name'].str.contains(<food_string>, case = False)), df['consumption'].str.extract(<creature_string>), np.nan)

but I must check for pairs so I tried to make a dictionary of food as keys and creatures as values , then make a string var of all the creatures for a given food key and look for those using str.contains :

unique_food = df2.food.unique()
food_dict = {elem : pd.DataFrame for elem in unique_food}
for key in food_dict.keys():
    food_dict[key] = df2[:][df2.food == key]

# create key:value pairs of food key and creature strings
food_strings = {}
for key, values in food_dict.items():
    food_strings.update({key: '|'.join(map(str, list(food_dict[key]['creature'].unique())))})

In[199]: food_strings
Out[199]: {'apple': 'squirrel|badger', 'banana': 'monkey|elephant'}

The problem is when I try to now apply str.contains:

for key, value in food_strings.items():
    np.where((df1['name'].str.contains('('+food_strings[key]+')', case = False)) & 
             (df1['consumption'].str.contains('('+food_strings[value]+')', case = False)), df1['consumptions'].str.extract('('+food_strings[value]+')'), np.nan)

I get a KeyError: .

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-62-7ab718066040> in <module>()
      1 for key, value in food_strings.items():
      2     np.where((df1['name'].str.contains('('+food_strings[key]+')', case = False)) & 
----> 3              (df1['consumption'].str.contains('('+food_strings[value]+')', case = False)), df1['consumption'].str.extract('('+food_strings[value]+')'), np.nan)

KeyError: 'squirrel|badger'

When I just try for only the value and not the key, it works for the first key:value pair but not the second:

for key in food_strings.keys():
    df1['test'] = np.where(df1['consumption'].str.contains('('+food_strings[key]+')', case =False), 
                                df1['consumption'].str.extract('('+food_strings[key]+')', expand=False), 
                                np.nan)

df1
Out[196]: 
            consumption                name      test
0     squirrelate apple               apple  squirrel
1     monkey likesapple          appleisred       NaN
2    monkey banana gets  banana is tropical       NaN
3     badger/getsbanana      banana is soft    badger
4    giraffe eats grass       lemon is sour       NaN
5    badger apple.loves   washington apples    badger
6      elephant is huge                kiwi       NaN
7        elephant/eats/             bananas       NaN
8  squirrel.digsingrass              apples  squirrel

I got the ones matching apple and squirrel|badger but missed banana:monkey|elephant.

can someone please help?


Solution

  • d1 = df1.dropna()
    d2 = df2.dropna()
    
    sump = d1.consumption.values.tolist()
    name = d1.name.values.tolist()
    cret = d2.creature.values.tolist() 
    food = d2.food.values.tolist()
    
    check = np.array(
        [
            [c in s and f in n for c, f in zip(cret, food)]
            for s, n in zip(sump, name)
        ]
    )
    
    # create a new series with the index of `d1` where we dropped na
    # then reindex with `df1.index` prior to `assign`
    test = pd.Series(check.dot(d2[['creature']].values).ravel(), d1.index)
    test = test.reindex(df1.index, fill_value='')
    df1.assign(test=test)
    
                consumption                name      test
    0     squirrelate apple               apple  squirrel
    1     monkey likesapple          appleisred          
    2    monkey banana gets  banana is tropical    monkey
    3     badger/getsbanana      banana is soft          
    4    giraffe eats grass       lemon is sour          
    5    badger apple.loves   washington apples    badger
    6      elephant is huge                kiwi          
    7        elephant/eats/             bananas  elephant
    8  squirrel.digsingrass              apples  squirrel