Добавление строк из одного CSV в другой, если их первый столбец совпадает - PullRequest
Новая
       5

Добавление строк из одного CSV в другой, если их первый столбец совпадает

0 голосов
/ 25 января 2019

У меня есть два файла csv с одинаковыми строками, но разными столбцами:

$ cat file1

category,a,b,c,d,e
apple,0,0,0,0,0
bear,1,1,1,1,1

$ cat file2

category,f,g,h,i,j
bear,10,10,10,10,10
apple,5,5,5,5,5

Ожидаемый результат: 

category,a,b,c,d,e,f,g,h,i,j
apple,0,0,0,0,0,5,5,5,5,5
bear,1,1,1,1,1,10,10,10,10,10

Оба file1 и file2 неупорядочены и имеют одинаковое количество совпадающих строк (~ 15000 строк).file1 имеет ~ 1000, а file2 имеет ~ 16000 столбцов.Я использовал следующий подход: 

import csv

with open ("file1.csv") as f:
    first = {rows[0]:rows[1:] for rows in list(csv.reader(f))}


with open("file2.csv") as f:
    for row in csv.reader(f):
        if row[0] in first:
            first[row[0]].extend(row[1:])

# print(first)
# {'category': ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j'], 'apple': ['0', '0', '0', '0', '0', '5', '5', '5', '5', '5'], 'bear': ['1', '1', '1', '1', '1', '10', '10','10', '10', '10']}

Я чувствую, что приближаюсь, но не могу написать first с тем же форматом, что и выше ожидаемый результат.

На заметку, имеет ли смысл сортировать файлы и добавлять их?У них одинаковое количество строк и категорий.

1 Ответ

0 голосов
/ 25 января 2019

Вы были довольно близки, но есть (IIUC) некоторые осложняющие факторы. Основным является то, что file_a и file_b имеют разную длину, и поэтому можно предположить, что они могут иметь несвязанное количество категорий в столбце А. То есть категории (bear / apple и т. Д.). ) может существовать в одном файле, а не в другом.

По этой причине я использую defaultdict. Это означает, что, если ключ не существует, он будет автоматически создан с пустым списком, сохраненным против него. Если мы видели этот ключ раньше, мы просто расширим существующий список. Если мы этого не сделаем, мы относимся к этому точно так же, но просто расширяем существующий пустой список по умолчанию .

Вы можете проверить это, поместив совершенно новую строку в file_a или file_b; код будет идти как обычно.

Кроме того, для упаковки в функцию (build_output) нам нужен только код with open()... один раз. Если вы хотите добавить другой файл, просто введите output = build_output('file_c.csv', output), и он будет добавлен таким же образом.

Наконец, мы используем первый столбец в качестве ключа в нашем словаре, а все остальные значения в файле как «значение». Нам нужно объединить их обратно в один список. Мы могли бы использовать первый столбец в качестве ключа, а затем сохранить всю строку против этого ключа, но это затруднило бы выписывание его обратно без многократного появления ключа. Вместо этого поместите key в список и объедините его со списком значений, чтобы получить один список, который мы можем записать в одну строку. 

import csv

from collections import defaultdict

output = defaultdict(list)


def build_output(file_name, output_obj):
    with open(file_name) as infile:
        reader_obj = csv.reader(infile)
        for row in reader_obj:
            output_obj[row[0]].extend(row[1:])
    return output_obj


def write_output(output):
    with open('output.csv', 'w', newline='') as outfile:
        writer = csv.writer(outfile)
        for key, value in output.items():
            row = [key] + value
            writer.writerow(row)


output = build_output('file_a.csv', output)
output = build_output('file_b.csv', output)
write_output(output)
...