如何修改Python代码从CSV按首列关键词提取指定区块数据

需要从CSV文件中，依据首列的Admissions或Readmissions关键词定位位置，向下移动3行、向右移动1列后，提取该位置开始的后续非空5列数据。当前基于re和pandas的Python代码会把两个关键词对应的目标数据区块合并输出，得修改代码实现按指定关键词单独输出对应的数据区块。

样本CSV数据 #

,Column1,Column2,Column3,Column4,Column5,Column6,Column7
Admissions,,,,,,,
,,,,,,,
,,,,,,,
,,,,,,,
,DataA1,DataA2,DataA3,DataA4,DataA5,,
,DataA6,DataA7,DataA8,DataA9,DataA10,,
,,,,,,,
Readmissions,,,,,,,
,,,,,,,
,,,,,,,
,,,,,,,
,DataR1,DataR2,DataR3,DataR4,DataR5,,
,DataR6,DataR7,DataR8,DataR9,DataR10,,

现有代码 #

import pandas as pd
import re

def extract_data(csv_path):
    df = pd.read_csv(csv_path, header=None)
    pattern = re.compile(r'Admissions|Readmissions')
    matches = df[0].str.contains(pattern, na=False)
    target_rows = []
    for idx in df[matches].index:
        start_row = idx + 3
        start_col = 1
        # 提取从start_row开始的非空行，取前5列
        data_block = df.loc[start_row:, start_col:start_col+4].dropna(how='all')
        target_rows.extend(data_block.values.tolist())
    return target_rows

# 调用示例
result = extract_data('sample.csv')
print(result)

错误输出 #

[['DataA1', 'DataA2', 'DataA3', 'DataA4', 'DataA5'], ['DataA6', 'DataA7', 'DataA8', 'DataA9', 'DataA10'], ['DataR1', 'DataR2', 'DataR3', 'DataR4', 'DataR5'], ['DataR6', 'DataR7', 'DataR8', 'DataR9', 'DataR10']]

当前输出把Admissions和Readmissions对应的数据合并成了一个列表，无法区分各自归属。

修正后的代码 #

import pandas as pd
import re

def extract_data_by_keyword(csv_path, target_keyword=None):
    df = pd.read_csv(csv_path, header=None)
    pattern = re.compile(r'Admissions|Readmissions')
    matches = df[0].str.contains(pattern, na=False)
    result_dict = {}
    
    for idx in df[matches].index:
        current_keyword = df.loc[idx, 0]
        # 指定关键词时，只处理匹配项
        if target_keyword and current_keyword != target_keyword:
            continue
            
        start_row = idx + 3
        start_col = 1
        # 提取从start_row开始的非空行，取前5列
        data_block = df.loc[start_row:, start_col:start_col+4].dropna(how='all')
        # 清洗数据，过滤空值行
        clean_data = [row.dropna().tolist() for _, row in data_block.iterrows() if any(row.notna())]
        result_dict[current_keyword] = clean_data
    
    # 按需返回结果：指定关键词则返回对应数据，否则返回全量字典
    return result_dict.get(target_keyword, []) if target_keyword else result_dict

# 调用示例
# 1. 单独获取Admissions对应的数据
admissions_data = extract_data_by_keyword('sample.csv', target_keyword='Admissions')
print("Admissions数据:", admissions_data)

# 2. 单独获取Readmissions对应的数据
readmissions_data = extract_data_by_keyword('sample.csv', target_keyword='Readmissions')
print("Readmissions数据:", readmissions_data)

# 3. 获取所有关键词对应的数据
all_data = extract_data_by_keyword('sample.csv')
print("所有数据:", all_data)

正确输出示例 #

Admissions数据: [['DataA1', 'DataA2', 'DataA3', 'DataA4', 'DataA5'], ['DataA6', 'DataA7', 'DataA8', 'DataA9', 'DataA10']]
Readmissions数据: [['DataR1', 'DataR2', 'DataR3', 'DataR4', 'DataR5'], ['DataR6', 'DataR7', 'DataR8', 'DataR9', 'DataR10']]
所有数据: {'Admissions': [['DataA1', 'DataA2', 'DataA3', 'DataA4', 'DataA5'], ['DataA6', 'DataA7', 'DataA8', 'DataA9', 'DataA10']], 'Readmissions': [['DataR1', 'DataR2', 'DataR3', 'DataR4', 'DataR5'], ['DataR6', 'DataR7', 'DataR8', 'DataR9', 'DataR10']]}

代码改动说明 #

1. 新增关键词过滤参数：添加target_keyword参数，传入具体关键词时只处理对应的数据区块，避免合并。
2. 用字典存储结果：把关键词和对应数据区块绑定，清晰区分不同来源的数据。
3. 增强数据清洗：额外过滤空值行，确保返回的都是有效非空数据。
4. 灵活调用逻辑：支持单独提取单个关键词数据，也可以一次性获取所有关键词的对应数据。

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