python - 如何将 csv 文件中的列中的数据拆分为两个单独的输出 csv 文件？

Question

我有一个 .csv 文件，例如：

ID NAME CATEGORIES
1, x,   AB
2, xx,  AA
3, xxx, BA

我将如何根据类别形成两个输出 .csv 文件，例如：

文件 1：

ID NAME CATEGORY
1, x,   A
2, xx,  A
3, xxx, B

文件 2：

ID NAME CATEGORY
1, x,   B
2, xx,  A
3, xxx, A

我已经设置了输入和输出，但只是一个空的 for 循环，我被难住了：

records = [line for line in csv.reader(open('test_input.csv', 'rt'), delimiter=',')]
outfile = open('test_output1.csv', 'wt')
outfileWriter = csv.writer(outfile, delimiter=',')

for record in records:
    #something!
    outfileWriter.writerow(record)

outfile.close()

我会很感激任何帮助！

Answer 1

import csv

records = [line for line in csv.reader(open('test_input.csv', 'rt'), delimiter=',')]
outfile1 = open('test_output1.csv', 'wt')
outfile2 = open('test_output2.csv', 'wt')
outfileWriter1 = csv.writer(outfile1, delimiter=',')
outfileWriter2 = csv.writer(outfile2, delimiter=',')

# headers always the same
outfileWriter1.writerow(records[0])
outfileWriter2.writerow(records[0])

for record in records[1:]:
    cat = record[-1].strip() # get category in form "AB"
    new_record = record
    new_record[-1] = "\t%s" % cat[0] # set category for file 1 with tab as a prefix
    outfileWriter1.writerow(new_record)
    new_record[-1] = "\t%s" % cat[1] # set category for file 2 with tab as a prefix
    outfileWriter2.writerow(new_record)

outfile1.close()
outfile2.close()

Answer 2

import csv
with open('input.csv') as f, open('file1.csv', 'w') as f1, open('file2.csv', 'w') as f2:
    header = next(f)    #read header
    reader = csv.reader(f, delimiter=',', skipinitialspace=True) 
    f1.write(header)    #write header
    f2.write(header)    #write header
    writ1 = csv.writer(f1, delimiter=',') 
    writ2 = csv.writer(f2, delimiter=',')
    for row in reader:
        c1, c2 = row[-1]          #split the category into c1 and c2
        writ1.writerow(row[:-1] + [c1])  #write c1 to file1
        writ2.writerow(row[:-1] + [c2])  #write c2 to file2