Добавить столбец в dataframe, используя регулярные выражения и словарь

Question

У меня есть данные как:

foo = pd.DataFrame({'id': ['A1', 'A2', 'A3', 'A4', 'A5', 'A6', 'A7', 'A8', 'A9', 'A10'], 
                    'amount': [10, 30, 40, 15, 20, 12, 55, 45, 60, 75], 
                    'description': [u'LYFT SAN FRANCISCO CA', u'XYZ STARBUCKS MINNEAPOLIS MN', u'HOLIDAY BEMIDJI MN', 
                                    u'MCDONALDS MADISON WI', u'ABC SUPERAMERICA MI', u'SUBWAY ROCHESTER MN', 
                                    u'NNT BURGER KING WI', u'UBER TRIP CA', u'superamerica CA', u'AMAZON NY']})

foo:

    id       amount description
    A1        10    LYFT SAN FRANCISCO CA
    A2        30    XYZ STARBUCKS MINNEAPOLIS MN
    A3        40    HOLIDAY BEMIDJI MN
    A4        15    MCDONALDS MADISON WI
    A5        20    ABC SUPERAMERICA MI
    A6        12    SUBWAY ROCHESTER MN
    A7        55    NNT BURGER KING WI
    A8        45    UBER TRIP CA
    A9        60    superamerica CA
    A10       75    AMAZON NY

Я хочу создать новый столбец, который классифицирует каждую запись на основе соответствия ключевого слова из descriptionстолбец.

Я использовал справку от этого ответа, чтобы сделать это следующим образом:

import re    
dict1 = {
    "LYFT" : "cab_ride",
    "UBER" : "cab_ride",
    "STARBUCKS" : "Food",
    "MCDONALDS" : "Food",
    "SUBWAY" : "Food",
    "BURGER KING" : "Food",
    "HOLIDAY" : "Gas",
    "SUPERAMERICA": "Gas"
        }

def get_category_from_desc(x):
    try:
        return next(dict1[k] for k in dict1 if re.search(k, x, re.IGNORECASE))
    except:
        return "Other"

foo['category'] = foo.description.map(get_category_from_desc)

Это работает, но я хочу спросить, является ли это лучшимвыход для этой проблемы.Поскольку у меня гораздо больший набор ключевых слов, которые могут указывать категорию, я должен создать огромный словарь:

dict1 = {
        "STARBUCKS" : "Food",
        "MCDONALDS" : "Food",
        "SUBWAY" : "Food",
        "BURGER KING" : "Food",
             .
             .
             .
        # ~50 more keys for "Food"

        "HOLIDAY" : "Gas",
        "SUPERAMERICA": "Gas",
             .
             .
             .
        # ~20 more keys for "Gas"

        "WALMART" : "grocery",
        "COSTCO": "grocery",
             .
             .
        # ..... ~30 more keys for "grocery"
             .
             .
        # ~ Many more categories with a large number of keys for each
}

Изменить : я также хочу знать, есть ли выходэто не требует от меня создания огромного словаря, подобного показанному выше.Могу ли я достичь этого с меньшей структурой данных, например:

dict2 = {
    "cab_ride" : ["LYFT", "UBER"], #....
    "food" : ["STARBUCKS", "MCDONALDS", "SUBWAY", "BURGER KING"], #....
    "gas" : ["HOLIDAY", "SUPERAMERICA"] #....
        }

Answer 1

Вы можете использовать .str метод доступа с extract и скомпилированное регулярное выражение, используя join для словарных ключей.

foo = pd.DataFrame({'id': ['A1', 'A2', 'A3', 'A4', 'A5', 'A6', 'A7', 'A8', 'A9', 'A10'], 
                    'amount': [10, 30, 40, 15, 20, 12, 55, 45, 60, 75], 
                    'description': [u'LYFT SAN FRANCISCO CA', u'XYZ STARBUCKS MINNEAPOLIS MN', u'HOLIDAY BEMIDJI MN', 
                                    u'MCDONALDS MADISON WI', u'ABC SUPERAMERICA MI', u'SUBWAY ROCHESTER MN', 
                                    u'NNT BURGER KING WI', u'UBER TRIP CA', u'superamerica CA', u'AMAZON NY']})


dict1 = {
    "LYFT" : "cab_ride",
    "UBER" : "cab_ride",
    "STARBUCKS" : "Food",
    "MCDONALDS" : "Food",
    "SUBWAY" : "Food",
    "BURGER KING" : "Food",
    "HOLIDAY" : "Gas",
    "SUPERAMERICA": "Gas"
        }

regstr = '(' + '|'.join(dict1.keys()) + ')'
foo['category'] = foo['description'].str.extract(regstr).squeeze().map(dict1).fillna('Other')
print(foo)

Выход:

    id  amount                   description  category
0   A1      10         LYFT SAN FRANCISCO CA  cab_ride
1   A2      30  XYZ STARBUCKS MINNEAPOLIS MN      Food
2   A3      40            HOLIDAY BEMIDJI MN       Gas
3   A4      15          MCDONALDS MADISON WI      Food
4   A5      20           ABC SUPERAMERICA MI       Gas
5   A6      12           SUBWAY ROCHESTER MN      Food
6   A7      55            NNT BURGER KING WI      Food
7   A8      45                  UBER TRIP CA  cab_ride
8   A9      60               superamerica CA     Other
9  A10      75                     AMAZON NY     Other

Answer 2

Я думаю, что это может быть достигнуто довольно легко, используя df.replace с заменой на основе регулярных выражений. Затем вы можете использовать df.where для обработки «других» дел.

dict2 = {rf'.*{k}.*': v for k, v in dict1.items()}

cats = foo['description'].replace(dict2, regex=True)
cats.where(cats != foo['description'], 'Other')

0    cab_ride
1        Food
2         Gas
3        Food
4         Gas
5        Food
6        Food
7    cab_ride
8       Other
9       Other
Name: description, dtype: object

---

Другой вариант - использование str.extract с map:

from collections import defaultdict

dict2 = defaultdict(lambda: 'Other')
dict2.update(dict1)

foo['description'].str.extract(rf"({'|'.join(dict1)})", expand=False).map(dict2)

0    cab_ride
1        Food
2         Gas
3        Food
4         Gas
5        Food
6        Food
7    cab_ride
8       Other
9       Other
Name: description, dtype: object