如何在Pandas中筛选仅含指定子串列表单个元素的行

需求说明 #

我有一个带编号和字符串列的Pandas DataFrame，要筛选出subjects列仅包含指定子串列表中单个元素的行——也就是说，该行的subjects里可以有其他不在指定列表中的内容，但来自指定列表的元素只能有一个。

原始数据 #

import pandas as pd

df = pd.DataFrame({
    "Value": ["P1", "P2", "P3", "P4", "P5", "P6", "P7"],
    "subjects": ["'A'", "'A','B'", "'A','B','C'", "'B'", "'C'", "'A','K'", "'B','ABC'"]
})

指定子串列表 #

sub = ["'A'","'B'","'C'"]

期望结果 #

Value    subjects
0    P1        "'A'"
3    P4        "'B'"
4    P5        "'C'"
5    P6    "'A','K'"
6    P7  "'B','ABC'"

现有方案问题 #

之前的实现通过多层循环拼接DataFrame，数据量大时性能极差，代码如下：

import pandas as pd
import numpy as np

completedf = pd.DataFrame()
for string in sub:
    for entry in df[df["subjects"].str.contains(str("'"+ string+"'")).replace(np.nan,False)]:
        for s in sub:
            if s != string:      
                m = df[~df["subjects"].str.contains(str("'"+ s[0]+"'")).replace(np.nan,False)]
    completedf = pd.concat([completedf, m])

高效解决方案 #

思路 #

统计每行subjects中属于指定子串列表的元素数量，只保留数量为1的行——这样既保证了有且仅有一个来自指定列表的元素，同时允许存在其他不在列表中的内容。

实现方式1：apply自定义函数（易读） #

def count_target_subjects(subjects_str, target_list):
    # 分割字符串为单个元素，处理掉引号和空格
    subjects = [s.strip().strip("'") for s in subjects_str.split(',')]
    # 转换为和target_list一致的带引号格式
    quoted_subjects = [f"'{s}'" for s in subjects]
    # 统计属于目标列表的元素数量
    return sum(1 for s in quoted_subjects if s in target_list)

# 计算每行的目标元素数量
df['target_count'] = df['subjects'].apply(count_target_subjects, target_list=sub)

# 筛选出目标元素数量为1的行，移除辅助列
result_df = df[df['target_count'] == 1].drop('target_count', axis=1)

print(result_df)

实现方式2：矢量化分组统计（性能更优） #

适合大数据量场景，用Pandas原生矢量化操作替代循环，性能提升明显：

# 分割subjects为多行，匹配目标列表后按Value分组计数
target_counts = df.assign(subject=df['subjects'].str.split(',')) \
                  .explode('subject') \
                  .assign(subject=lambda x: x['subject'].str.strip()) \
                  .query('subject in @sub') \
                  .groupby('Value')['subject'] \
                  .count()

# 筛选出计数为1的行
result_df = df[df['Value'].isin(target_counts[target_counts == 1].index)]

print(result_df)

为什么高效 #

• 避免了循环拼接DataFrame的操作，减少不必要的内存占用和重复计算
• 矢量化操作利用Pandas底层优化，比Python循环快几个数量级，尤其适合大数据量场景

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