如何高效过滤Pandas DataFrame中TCK列含不同条目的ID？

问题描述 #

现有一个Pandas DataFrame，需筛选出所有TCK列（每个元素是包含逗号分隔字符串的列表）中条目不全相同的ID。示例数据及期望输出如下：

示例输入 #

import pandas as pd
import numpy as np

df1 = pd.DataFrame({"ID": [1, 2, 3, 4],
        "TCK": [["AA, AA, AC"], ["LL, LL"], ["DD , DB, DF, DE"], ["LO , LO, LO, LO, LO, LO"]]})

期望输出 #

df2 = pd.DataFrame({"ID": [1, 3],
        "TCK": [["AA, AA, AC"],["DD , DB, DF, DE"]]})

同时需要处理np.nan的情况，避免循环带来的复杂度，寻求高效优雅的解决方案。

解决方案 #

通过apply结合集合去重的方式快速判断，同时兼容空值处理：

import pandas as pd
import numpy as np

# 定义筛选掩码
mask = df1['TCK'].apply(
    lambda x: len(set(s.strip() for s in x[0].split(',')) - {''}) > 1 
    if not pd.isna(x) else False
)

# 筛选并重置索引
df2 = df1[mask].reset_index(drop=True)

代码解释 #

1. 空值处理：先判断TCK元素是否为np.nan，若是则标记为不保留（返回False）
2. 字符串拆分与清洗：取出列表中的逗号分隔字符串，按逗号拆分后去除每个元素的首尾空格
3. 唯一性判断：将清洗后的元素转为集合，去掉空字符串（处理多余逗号导致的无效空值），若集合长度大于1，说明存在不同条目，标记为保留（返回True）
4. 筛选结果：用掩码筛选原DataFrame，重置索引得到最终结果

兼容空值的测试示例 #

如果原DataFrame包含np.nan：

df_with_nan = pd.DataFrame({"ID": [1,2,3,4,5],
        "TCK": [["AA, AA, AC"], ["LL, LL"], np.nan, ["DD , DB"], ["LO, LO, LP"]]})

运行上述代码后，会自动跳过ID为3的行，保留ID为1、4、5的行。

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