Python按指定列表排序表格行并保留同ID所有行的问题求助

需求是按给定的ID列表顺序，对第一列为ID的表格行排序，同时保留每个ID对应的所有行，但当前脚本仅返回每个ID的第一行，需要修复。

示例数据 #

原始表格（table_data.txt）

A       B
1211    ds
3245    ssssd
3114    dsf
3114    apple
4324    sssvdff
4324    weewr
4324    bla 
4324    orange
1211    something
1211    blue  

指定ID顺序（Lorder.txt）

4324
3245
3114
1211

当前脚本输出

A       B
4324    sssvdff
3245    ssssd
3114    dsf
1211    something

期望输出

A       B
4324    sssvdff
4324    weewr
4324    bla 
4324    orange
3245    ssssd
3114    dsf
3114    apple
1211    something
1211    blue  

现有代码 #

import re

def map_ids_to_row_list(_id):
    _id = _id.strip('\n')
    for s in lines:
        if re.search(_id,s):
            return s

file = open('Lorder.txt', 'r').readlines()
names = []
for name in file:
    names.append(name.strip('\n'))

name = 'table_data.txt'
name_split = name.rsplit('.',1)
new_name = name_split[0] + '_sorted.csv'
    
table = open(name,'r')
table_read = table.readlines()
    
lines = []

for s in table_read:
    lines.append(s)
    
ordered_table = list(map(map_ids_to_row_list, names))
    
with open(new_name,'w') as fs:
    for row in ordered_table:
        fs.write(str(row))

问题原因 #

原代码的map_ids_to_row_list函数在找到第一个匹配ID的行后就直接return，导致每个ID仅能获取到第一行数据，无法保留所有行。

修复方案 #

核心思路是先将所有表格行按ID分组存储，再按照指定ID顺序依次取出对应分组的所有行：

# 读取指定ID顺序
with open('Lorder.txt', 'r') as f:
    id_order = [line.strip() for line in f if line.strip()]

# 读取表格数据，分离表头和内容行
with open('table_data.txt', 'r') as f:
    table_lines = [line.rstrip('\n') for line in f]
header = table_lines[0]
data_lines = table_lines[1:]

# 按ID分组存储所有行
id_to_rows = {}
for line in data_lines:
    # 分割第一列作为ID（假设ID和内容用空格/制表符分隔）
    row_id = line.split()[0]
    if row_id not in id_to_rows:
        id_to_rows[row_id] = []
    id_to_rows[row_id].append(line)

# 按指定ID顺序生成排序后的表格
sorted_table = [header]
for _id in id_order:
    if _id in id_to_rows:
        sorted_table.extend(id_to_rows[_id])

# 写入结果文件
name = 'table_data.txt'
new_name = name.rsplit('.',1)[0] + '_sorted.csv'
with open(new_name, 'w') as fs:
    fs.write('\n'.join(sorted_table) + '\n')

说明 #

1. 用with语句管理文件，避免资源泄漏
2. 先将数据按ID分组到字典中，确保每个ID的所有行都被保存
3. 按照指定ID顺序遍历，将对应分组的所有行追加到结果列表中
4. 保留原表格的表头，保证输出格式一致

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