Search code examples
pythonpandasdictionarydataframeapply

map str.contains across pandas DataFrame


Beginner with python - I'm looking to create a dictionary mapping of strings, and the associated value. I have a dataframe and would like create a new column where if the string matches, it tags the column as x.

df = pd.DataFrame({'comp':['dell notebook', 'dell notebook S3', 'dell notepad', 'apple ipad', 'apple ipad2', 'acer chromebook', 'acer chromebookx', 'mac air', 'mac pro', 'lenovo x4'],
              'price':range(10)})

For Example I would like to take the above df and create a new column df['company'] and set it to a mapping of strings.

I was thinking of doing something like

product_map = {'dell':'Dell Inc.',
               'apple':'Apple Inc.',
               'acer': 'Acer Inc.',
               'mac': 'Apple Inc.',
               'lenovo': 'Dell Inc.'}

Then I wanted to iterate through it to check the df.comp column and see if each entry contained one of those strings, and to set the df.company column to the value in the dictionary.

Not sure how to do this correctly though.


Solution

  • There are many ways to do this. One way to do it would be the following:

    def like_function(x):
        group = "unknown"
        for key in product_map:
            if key in x:
                group = product_map[key]
                break
        return group
    
    df['company'] = df.comp.apply(like_function)