Searching multiple substrings in a column of strings and return substring category

user_5141727 0
Viewed 367

I have two dataframes as follows:

df1 = pd.DataFrame({"id":["01", "02", "03", "04", "05", "06"],
                    "string":["This is a cat",
                              "That is a dog",
                              "Those are birds",
                              "These are bats",
                              "I drink coffee",
                              "I bought tea"]})

df2 = pd.DataFrame({"category":[1, 1, 2, 2, 3, 3],
                    "keywords":["cat", "dog", "birds", "bats", "coffee", "tea"]})

My dataframes look like this

df1:

id   string
01   This is a cat
02   That is a dog
03   Those are birds
04   These are bats
05   I drink coffee
06   I bought tea

df2:

category   keywords
1          cat
1          dog
2          birds
2          bats
3          coffee
3          tea

I would like to have an output column on df1 which is the category if at least one keyword in df2 is detected in each string in df1, otherwise return None. The expected output should be the following.

id   string             category
01   This is a cat         1
02   That is a dog         1
03   Those are birds       2
04   These are bats        2
05   I drink coffee        3
06   I bought tea          3

I can think of looping over keywords one-by-one and scan through string one-by-one but it is not efficient enough if data is getting big. May I have your suggestions how to improve? Thank you.

pandas dataframe lookup python string
3 Answers
# Modified your data a bit.
df1 = pd.DataFrame({"id":["01", "02", "03", "04", "05", "06", "07"],
                    "string":["This is a cat",
                              "That is a dog",
                              "Those are birds",
                              "These are bats",
                              "I drink coffee",
                              "I bought tea", 
                              "This won't match squat"]})

You can use a list comprehension involving next with the default argument.

df1['category'] = [
    next((c for c, k in df2.values if k in s), None) for s in df1['string']] 

df1
   id                  string  category
0  01           This is a cat       1.0
1  02           That is a dog       1.0
2  03         Those are birds       2.0
3  04          These are bats       2.0
4  05          I drink coffee       3.0
5  06            I bought tea       3.0
6  07  This won't match squat       NaN

You cannot avoid the O(N2) complexity, but this should be quite performant since it doesn't always have to iterate over every string in the inner loop (unless in the worst case).

Note that this currently only supports substring matches (not regex-based matching, although with a little modification that can be done).

user_49090870

Use list comprehension with split and match by dictionary created by df2:

d = dict(zip(df2['keywords'], df2['category']))
df1['cat'] = [next((d[y] for y in x.split() if y in d), None) for x in df1['string']]

print (df1)
   id           string  cat
0  01    This is a cat  1.0
1  02    That is a dog  1.0
2  03  Those are birds  2.0
3  04   These are bats  2.0
4  05   I drink coffee  3.0
5  06    I bought thea  NaN
user_29010020

Another easy to understand solution mapping df1['string']:

# create a dictionary with keyword->category pairs
cats = dict(zip(df2.keywords, df2.category))

def categorize(s):
    for cat in cats.keys():
        if cat in s:
            return cats[cat]
    # return 0 in case nothing is found
    return 0

df1['category'] = df1['string'].map(lambda x: categorize(x))

print(df1)

   id           string  category
0  01    This is a cat         1
1  02    That is a dog         1
2  03  Those are birds         2
3  04   These are bats         2
4  05   I drink coffee         3
5  06     I bought tea         3
user_73907200
Related