Python:尝试从一个csv文件提取行并将其输出到另一个csv文件

import csv

def extract_rows(input_file, output_file, rows):
    with open(input_file, 'r') as file:
        reader = csv.reader(file)
        data = list(reader)
    
    extracted_data = [data[row] for row in rows]
    
    with open(output_file, 'w', newline='') as file:
        writer = csv.writer(file)
        writer.writerows(extracted_data)

# 示例用法
input_file = 'input.csv'
output_file = 'output.csv'
rows_to_extract = [1, 3, 5]  # 假设要提取第1、3、5行

extract_rows(input_file, output_file, rows_to_extract)