使用正则表达式和字典将列添加到数据框答案

【问题标题】：Add column to dataframe using regex and dictionary使用正则表达式和字典将列添加到数据框
【发布时间】：2019-04-18 19:37:44
【问题描述】：

我有这样的数据：

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']})

富：

    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 列的关键字匹配对每条记录进行分类。

我已使用this 的帮助，通过以下方式进行操作：

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"] #....
        }

【问题讨论】：

至于您的编辑：可能没有使用看起来像那样的字典...
那么字典应该是什么样子？
它必须是平的，不幸的是...dict3 = {v: k for k, V in dict2.items() for v in V}

标签： python regex pandas

【解决方案1】：

我认为这可以很容易地使用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

【讨论】：

rf... 带有 f 字符串格式的漂亮正则表达式。谢谢你教我一些新东西！ +1
@ScottBoston 我只希望 OP 有 python3.6...在这里赌一把：P

【解决方案2】：

您可以将.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

【讨论】：