基于Pandas的逐行复杂数据子集提取技术需求

问题背景 #

需用Pandas导入CSV文件后提取数据子集，核心难点是在单元格的数组元素中查找'a'，并对应提取关联的smd值，同时要兼容未来新增同类列的需求。

示例数据 #

import pandas as pd

data = {'out_type_1': [['a', 'Reading primary outcome'], ['a', 'Other outcome'], ['a', 'Reading primary outcome'], 
                       ['a', 'Reading primary outcome'], ['a', 'Other outcome'], ['a', 'Other outcome'], 
                       ['a', 'Other outcome'], ['a', 'Other outcome'], ['Other outcome'], ['Other outcome'], 
                       ['a', 'Other outcome'], ['a', 'Reading primary outcome'], ['a', 'Reading primary outcome'], 
                       ['a', 'Other outcome'], ['a', 'Reading primary outcome'], ['Reading primary outcome'], 
                       ['a', 'Reading primary outcome'], None, None, ['a', 'Other outcome'], ['a', 'Reading primary outcome'], 
                       ['a', 'Reading primary outcome'], ['a', 'Other outcome']], 
        'smd_1': [-0.045, 0.0478, 0.1843, 0.0534, -0.0465, 0.2039, 0.6767, 0.5996, 0, 0.0517, 0.0922, 0.0631, 0.184, 0.1258, 0.2, -0.1982, 0.5207, 0.0187, 0.1245, 0.9315, 0.3784, 0.5811, 0.3998], 
        'out_type_2': [['Other outcome'], ['Other outcome'], ['Other outcome'], ['Other outcome'], None, None, None, None, 
                       ['a', 'Reading primary outcome'], ['a', 'Reading primary outcome'], ['Other outcome'], 
                       ['Mathematics primary outcome'], ['Other outcome'], ['Mathematics primary outcome'], 
                       ['Mathematics primary outcome'], ['a', 'Mathematics primary outcome'], None, None, None, None, 
                       ['Mathematics primary outcome'], ['Other outcome'], ['Other outcome']], 
        'smd_2': [0.9197, 0.1541, 0.1229, -0.1277, None, None, None, None, -0.1, -0.04, 0.5401, 0.174, 0.2519, 0.0193, 0.49, -0.0209, None, 0.1801, 0.0199, None, 0.5347, -0.1127, 0.3197]}

df = pd.DataFrame(data)

需求说明 #

• 逐行处理数据，将包含'a'的out_type列值及其对应的smd列值存入新列out_type_subset和smd_subset
• 每行最多存在一个含'a'的out_type列，无匹配项时新列保留NA
• 兼容未来新增的out_type_n和smd_n列

期望输出 #

import pandas as pd

data = {
    'out_type_subset': [
        "['a', 'Reading primary outcome']",
        "['a', 'Other outcome']",
        "['a', 'Reading primary outcome']",
        "['a', 'Reading primary outcome']",
        "['a', 'Other outcome']",
        "['a', 'Other outcome']",
        "['a', 'Other outcome']",
        "['a', 'Other outcome']",
        "['a', 'Reading primary outcome']",
        "['a', 'Reading primary outcome']",
        "['a', 'Other outcome']",
        "['a', 'Reading primary outcome']",
        "['a', 'Reading primary outcome']",
        "['a', 'Other outcome']",
        "['a', 'Reading primary outcome']",
        "['a', 'Mathematics primary outcome']",
        "['a', 'Reading primary outcome']",
        pd.NA,
        pd.NA,
        "['a', 'Other outcome']",
        "['a', 'Reading primary outcome']",
        "['a', 'Reading primary outcome']",
        "['a', 'Other outcome']"
    ],
    'smd_subset': [
        -0.045,
        0.0478,
        0.1843,
        0.0534,
        -0.0465,
        0.2039,
        0.6767,
        0.5996,
        -0.1,
        -0.04,
        0.0922,
        0.0631,
        0.184,
        0.1258,
        0.2,
        -0.0209,
        0.5207,
        pd.NA,
        pd.NA,
        0.9315,
        0.3784,
        0.5811,
        0.3998
    ]
}

df = pd.DataFrame(data)

解决方案 #

import pandas as pd

# 若从CSV导入，替换为 df = pd.read_csv('your_file.csv')
data = {'out_type_1': [['a', 'Reading primary outcome'], ['a', 'Other outcome'], ['a', 'Reading primary outcome'], 
                       ['a', 'Reading primary outcome'], ['a', 'Other outcome'], ['a', 'Other outcome'], 
                       ['a', 'Other outcome'], ['a', 'Other outcome'], ['Other outcome'], ['Other outcome'], 
                       ['a', 'Other outcome'], ['a', 'Reading primary outcome'], ['a', 'Reading primary outcome'], 
                       ['a', 'Other outcome'], ['a', 'Reading primary outcome'], ['Reading primary outcome'], 
                       ['a', 'Reading primary outcome'], None, None, ['a', 'Other outcome'], ['a', 'Reading primary outcome'], 
                       ['a', 'Reading primary outcome'], ['a', 'Other outcome']], 
        'smd_1': [-0.045, 0.0478, 0.1843, 0.0534, -0.0465, 0.2039, 0.6767, 0.5996, 0, 0.0517, 0.0922, 0.0631, 0.184, 0.1258, 0.2, -0.1982, 0.5207, 0.0187, 0.1245, 0.9315, 0.3784, 0.5811, 0.3998], 
        'out_type_2': [['Other outcome'], ['Other outcome'], ['Other outcome'], ['Other outcome'], None, None, None, None, 
                       ['a', 'Reading primary outcome'], ['a', 'Reading primary outcome'], ['Other outcome'], 
                       ['Mathematics primary outcome'], ['Other outcome'], ['Mathematics primary outcome'], 
                       ['Mathematics primary outcome'], ['a', 'Mathematics primary outcome'], None, None, None, None, 
                       ['Mathematics primary outcome'], ['Other outcome'], ['Other outcome']], 
        'smd_2': [0.9197, 0.1541, 0.1229, -0.1277, None, None, None, None, -0.1, -0.04, 0.5401, 0.174, 0.2519, 0.0193, 0.49, -0.0209, None, 0.1801, 0.0199, None, 0.5347, -0.1127, 0.3197]}

df = pd.DataFrame(data)

# 自动识别所有out_type列及对应的smd列，兼容新增列
out_type_cols = [col for col in df.columns if col.startswith('out_type_')]
smd_cols = [col.replace('out_type', 'smd') for col in out_type_cols]

# 定义行处理函数
def extract_matching_row(row):
    for out_col, smd_col in zip(out_type_cols, smd_cols):
        out_val = row[out_col]
        # 检查值非空、为列表且包含'a'
        if out_val is not None and isinstance(out_val, list) and 'a' in out_val:
            return pd.Series([str(out_val), row[smd_col]])
    # 无匹配项返回NA
    return pd.Series([pd.NA, pd.NA])

# 应用函数生成新列
df[['out_type_subset', 'smd_subset']] = df.apply(extract_matching_row, axis=1)

# 生成结果DataFrame（可选保留原列或只保留新列）
result_df = df[['out_type_subset', 'smd_subset']]
print(result_df)

# 导出到CSV（如果需要）
# result_df.to_csv('subset_result.csv', index=False)

代码逻辑说明 #

1. 自动列识别：通过列名前缀匹配所有out_type_*列，并对应生成smd_*列列表，无需手动指定列名，兼容未来新增列
2. 行遍历检查：逐行检查每对out_type和smd列，判断out_type的值是否为非空列表且包含'a'
3. 结果返回：找到匹配项时返回out_type的字符串形式和对应smd值；无匹配项返回NA
4. 导出支持：最后可将结果导出为CSV文件

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