如何使用contains运算符为Pandas DataFrame新增两列？

我从Excel文件导入了一个记录不同类型支出的DataFrame，该DataFrame包含名为CONCEPT的列。我希望根据支出类型新增EXPENSE TYPE和SUBCATEGORY两列，通过识别CONCEPT中的关键词来匹配对应值。现有代码仅能实现新增一列EXPENSE TYPE：

import pandas as pd

dictionary = {
      "DATE" : ['12/02/2023', '02/01/2023', '02/01/2023', '10/02/2023'],    
      "CONCEPT" : ['Supermarket','Restaurant', 'petrol station', 'decathlon'],
      "EUR" : [-150,-50,-45,-95]
          }
df = pd.DataFrame(dictionary)

df['EXPENSE TYPE'] = pd.Series(dtype="string")

df['EXPENSE TYPE'][df['CONCEPT'].str.upper().str.contains('SUPERMARKET')] = 'FOOD'
df['EXPENSE TYPE'][df['CONCEPT'].str.upper().str.contains('RESTAURANT')] = 'FOOD'
df['EXPENSE TYPE'][df['CONCEPT'].str.upper().str.contains('PETROL')] = 'GAS'
df['EXPENSE TYPE'][df['CONCEPT'].str.upper().str.contains('DECATHLON')] = 'CLOTHES' 

运行上述代码得到的结果：

DATE         CONCEPT  EUR EXPENSE TYPE
0  12/02/2023     Supermarket -150         FOOD
1  02/01/2023      Restaurant  -50         FOOD
2  02/01/2023  petrol station  -45          GAS
3  10/02/2023       decathlon  -95      CLOTHES

期望输出：

DATE         CONCEPT  EUR EXPENSE TYPE  SUBCATEGORY
0  12/02/2023     Supermarket -150         FOOD         HOME
1  02/01/2023      Restaurant  -50         FOOD  OUT OF HOME
2  02/01/2023  petrol station  -45          GAS       DIESEL
3  10/02/2023       decathlon  -95      CLOTHES        SPORT

请问如何调整代码实现同时新增这两列？

方法一：延续原逻辑，批量赋值两列 #

先初始化目标列，再通过loc索引一次性给匹配行的两列赋值，避免链式索引引发的警告：

import pandas as pd

dictionary = {
      "DATE" : ['12/02/2023', '02/01/2023', '02/01/2023', '10/02/2023'],    
      "CONCEPT" : ['Supermarket','Restaurant', 'petrol station', 'decathlon'],
      "EUR" : [-150,-50,-45,-95]
          }
df = pd.DataFrame(dictionary)

# 初始化两列为string类型
df['EXPENSE TYPE'] = pd.Series(dtype="string")
df['SUBCATEGORY'] = pd.Series(dtype="string")

# 匹配关键词并批量设置列值
mask_super = df['CONCEPT'].str.upper().str.contains('SUPERMARKET')
df.loc[mask_super, ['EXPENSE TYPE', 'SUBCATEGORY']] = ['FOOD', 'HOME']

mask_resto = df['CONCEPT'].str.upper().str.contains('RESTAURANT')
df.loc[mask_resto, ['EXPENSE TYPE', 'SUBCATEGORY']] = ['FOOD', 'OUT OF HOME']

mask_petrol = df['CONCEPT'].str.upper().str.contains('PETROL')
df.loc[mask_petrol, ['EXPENSE TYPE', 'SUBCATEGORY']] = ['GAS', 'DIESEL']

mask_decathlon = df['CONCEPT'].str.upper().str.contains('DECATHLON')
df.loc[mask_decathlon, ['EXPENSE TYPE', 'SUBCATEGORY']] = ['CLOTHES', 'SPORT']

方法二：用规则字典+np.select（更易维护） #

如果后续关键词或分类规则增多，用集中管理的规则字典会更高效：

import pandas as pd
import numpy as np

dictionary = {
      "DATE" : ['12/02/2023', '02/01/2023', '02/01/2023', '10/02/2023'],    
      "CONCEPT" : ['Supermarket','Restaurant', 'petrol station', 'decathlon'],
      "EUR" : [-150,-50,-45,-95]
          }
df = pd.DataFrame(dictionary)

# 定义匹配规则：(匹配条件, (EXPENSE TYPE值, SUBCATEGORY值))
rules = [
    (df['CONCEPT'].str.upper().str.contains('SUPERMARKET'), ('FOOD', 'HOME')),
    (df['CONCEPT'].str.upper().str.contains('RESTAURANT'), ('FOOD', 'OUT OF HOME')),
    (df['CONCEPT'].str.upper().str.contains('PETROL'), ('GAS', 'DIESEL')),
    (df['CONCEPT'].str.upper().str.contains('DECATHLON'), ('CLOTHES', 'SPORT'))
]

# 拆分条件与对应值列表
conditions = [cond for cond, _ in rules]
expense_type_vals = [val[0] for _, val in rules]
subcategory_vals = [val[1] for _, val in rules]

# 生成目标列
df['EXPENSE TYPE'] = np.select(conditions, expense_type_vals, default=np.nan).astype('string')
df['SUBCATEGORY'] = np.select(conditions, subcategory_vals, default=np.nan).astype('string')

内容的提问来源于stack exchange，提问作者nekovolta