从DataFrame提取指定子串并写入新DataFrame的实现方法

以下几种方式可以实现你的需求：

方法一：使用np.select（逻辑清晰，推荐） #

适合多条件分支场景，明确指定每个条件对应的值：

import numpy as np

# 定义匹配条件与对应输出值
conditions = [
    df_input['Visit  '].str.contains('Pre'),
    df_input['Visit  '].str.contains('Post'),
    df_input['Visit  '].str.contains('Screening')
]
choices = ['Pre', 'Post', 'Screening']

# 赋值到目标列，无匹配项时返回NaN，可修改default参数设置默认值
df_output['VISIT'] = np.select(conditions, choices, default=np.nan)

方法二：使用str.extract直接提取关键词 #

通过正则表达式精准匹配并提取目标字符串：

# 正则匹配三个关键词中的任意一个，无匹配返回NaN
df_output['VISIT'] = df_input['Visit  '].str.extract(r'(Pre|Post|Screening)', expand=False)

方法三：自定义函数+apply #

如果需要更灵活的逻辑判断，可自定义函数处理：

def get_visit_type(visit_text):
    if 'Pre' in visit_text:
        return 'Pre'
    elif 'Post' in visit_text:
        return 'Post'
    elif 'Screening' in visit_text:
        return 'Screening'
    else:
        return np.nan  # 可替换为其他默认值，比如空字符串''

df_output['VISIT'] = df_input['Visit  '].apply(get_visit_type)

之前代码的问题说明 #

• 第一段代码：df_input[df_input['Visit '].str.contains('Pr|Po|Sc')]返回的是筛选后的整行数据，而df_output['VISIT']是单列，两者长度不匹配，因此报错"Columns must be the same length as key"。
• 第二段代码：str.contains本身返回布尔值（True/False），自然会将布尔值写入目标列，不符合提取关键词的需求。

内容的提问来源于stack exchange，提问作者A-Simone