嵌套循环按键合并两个CSV文件仅首次迭代生效问题

问题根源 #

• csv.reader是一次性迭代器：第一次遍历f2_reader后，文件指针已移到file2末尾，后续外层循环再进入内层循环时，f2_reader无剩余数据可遍历，导致程序仅处理file1第一行就停止。
• 冗余操作：循环内重复创建csv.writer(f3)，完全没必要，初始化一次即可复用。
• 效率低下：嵌套循环时间复杂度为O(n²)，数据量大时运行极慢。

修复方案 #

优化思路 #

1. 提前将file2数据加载到字典中，以共同键（file2第一列）为key、行数据为value，后续匹配时可直接O(1)查找，彻底解决迭代器耗尽问题。
2. 移除冗余的writer创建操作，复用初始化好的writer对象。
3. 复制原行数据再修改，避免意外影响原始数据。

修复后的代码 #

import csv
import os

output_dir = "<output_dir>"
file1 = "f1.csv"
file2 = "f2.csv"

path1 = os.path.join(output_dir, file1)
path2 = os.path.join(output_dir, file2)

file3 = "fd_labels_processed.csv"
output_file = os.path.join(output_dir, file3)

# 预加载file2数据到字典，用共同键作为索引
f2_data = {}
with open(path2, 'r') as f2:
    f2_reader = csv.reader(f2, delimiter=',')
    next(f2_reader)  # 跳过表头
    for row_f2 in f2_reader:
        # 若存在重复键，可将value改为列表存储多行，此处假设键唯一
        f2_data[row_f2[0]] = row_f2

with open(path1, 'r') as f1, open(output_file, 'w+', newline='') as f3:
    f1_reader = csv.reader(f1, delimiter=',')
    next(f1_reader)  # 跳过表头
    writer = csv.writer(f3, delimiter=',')
    
    # 写入新文件表头
    writer.writerow(["ATTRIBUTE_1", "ATTRIBUTE_2", "ATTRIBUTE_3", "ATTRIBUTE_4", "ATTRIBUTE_5", "ATTRIBUTE_6", "ATTRIBUTE_7", "ATTRIBUTE_8", "ATTRIBUTE_9"])
    
    count = 0
    for row_f1 in f1_reader:
        key = row_f1[1]  # file1的共同键为第二列
        if key in f2_data:
            row_f2 = f2_data[key]
            # 匹配额外条件
            if (row_f1[3] == row_f2[2] and row_f1[4] == row_f2[3] and 
                row_f1[5] == row_f2[4] and row_f1[6] == row_f2[5]):
                # 复制行数据再修改，避免影响字典中的原始数据
                new_row = row_f2.copy()
                new_row.append(row_f1[9])
                new_row.append(row_f1[11])
                new_row.append(row_f1[12])
                writer.writerow(new_row)
                count += 1
                print(count)

额外说明 #

• 若file2存在重复的共同键，需将字典的value改为列表，存储所有对应行后再遍历匹配条件。
• 移除了未使用的numpy导入，精简代码。
• 用writerow替代writerows([row_f2])，单行写入更简洁。

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