Слияние данных на основе сопоставления первого столбца в Python

В настоящее время у меня есть два набора файлов данных, которые выглядят так:

Файл 1:

test1 ba ab cd dh gf
test2 fa ab cd dh gf
test3 rt ty er wq ee
test4 er rt sf sd sa

и в файле 2:

test1 123 344 123
test1 234 567 787
test1 221 344 566
test3 456 121 677

Я хотел бы объединить файлы, основанные на строках в первом столбце (чтобы "тесты" совпадали)

вот так:

test1 ba ab cd dh gf 123 344 123
test1 ba ab cd dh gf 234 567 787
test1 ba ab cd dh gf 221 344 566
test3 rt ty er wq ee 456 121 677

У меня есть этот код

def combineFiles(file1,file2,outfile):
      def read_file(file):
         data = {}
         for line in csv.reader(file):
            data[line[0]] = line[1:]
         return data
      with open(file1, 'r') as f1, open(file2, 'r') as f2:
         data1 = read_file(f1)
         data2 = read_file(f2)
         with open(outfile, 'w') as out:
            wtr= csv.writer(out)
            for key in data1.keys():
               try:
                  wtr.writerow(((key), ','.join(data1[key]), ','.join(data2[key])))
               except KeyError:
                  pass

Однако результат в конечном итоге выглядит следующим образом:

test1 ba ab cd dh gf 123 344 123
test3 er rt sf sd sa 456 121 677

Может кто-нибудь помочь мне с тем, как сделать вывод, чтобы test1 можно было распечатать все три раза?

Очень признателен

Ответы

Ответ 1

Хотя я бы порекомендовал подход Брэда Соломона, так как он довольно лаконичен, вам просто нужно внести небольшие изменения в ваш код.

Поскольку ваш второй файл имеет последнее слово, вам просто нужно создать словарь для первого файла. Затем вы можете записать выходной файл, как вы читаете из второго файла, выбирая значения из словаря data1 по ходу дела:

with open(file1, 'r') as f1, open(file2, 'r') as f2:
    data1 = read_file(f1)
    with open(outfile, 'w') as out:
        wtr = csv.writer(out, delimiter=' ')
        for line in csv.reader(f2, delimiter=' '):
            # only write if there is a corresponding line in file1
            if line[0] in data1:
                # as you write, get the corresponding file1 data
                wtr.writerow(line[0:] + data1[line[0]] + line[1:])

Ответ 2

Вы могли бы хотеть, чтобы дать панду библиотеке попробовать; это облегчает такие случаи:

>>> import pandas as pd
>>> pd.merge(df1, df2, on='testnum', how='inner')
  testnum 1_x 2_x 3_x   4   5  1_y  2_y  3_y
0   test1  ba  ab  cd  dh  gf  123  344  123
1   test1  ba  ab  cd  dh  gf  234  567  787
2   test1  ba  ab  cd  dh  gf  221  344  566
3   test3  rt  ty  er  wq  ee  456  121  677

Предполагается, что столбец test называется "testnum".

>>> df1
  testnum   1   2   3   4   5
0   test1  ba  ab  cd  dh  gf
1   test2  fa  ab  cd  dh  gf
2   test3  rt  ty  er  wq  ee
3   test4  er  rt  sf  sd  sa

>>> df2
  testnum    1    2    3
0   test1  123  344  123
1   test1  234  567  787
2   test1  221  344  566
3   test3  456  121  677

Вы прочитали бы это с помощью pd.read_csv().

Ответ 3

Проблема в том, что вы перезаписываете ключи в строке

data[line[0]] = line[1:]

Поскольку ваши файлы имеют неуникальные "ключи", вы можете попробовать вручную сделать их уникальными, используя enumerate:

for ind, line in enumerate(csv.reader(file)):
    unique_key = ''.join([line[0], "_", str(ind)])
    data[unique_key] = line[1:]

Позже, когда вы объединяете результат, вы можете удалить ключи после подчеркивания:

wtr.writerow(((key.split("_")[0], ','.join(data1[key]), ','.join(data2[key])))

На мой вкус, все это очень неуклюже. Если ваша цель - читать, манипулировать и записывать данные в и из CSV файлов, я бы порекомендовал изучить панды, так как этот код можно написать в несколько строк с использованием DataFrames (см. Ответ Брэда Соломона).

Ответ 4

Вы можете попытаться собрать свои элементы в отдельную collections.defaultdict(), а затем получить декартово произведение пересекающихся строк с помощью itertools.product():

from collections import defaultdict
from itertools import product

def collect_rows(file):
    collection = defaultdict(list)

    for line in file:
        col1, *rest = line.split()
        collection[col1].append(rest)

    return collection

with open("file1.txt") as f1, open("file2.txt") as f2, open("output.txt", "w") as out:
    f1_collection = collect_rows(f1)
    f2_collection = collect_rows(f2)

    # Ordered intersection, no need to sort
    set_2 = set(f2_collection)
    intersection = [key for key in f1_collection if key in set_2]

    for key in intersection:
        for x, y in product(f1_collection[key], f2_collection[key]):
            out.write("%s\n" % " ".join([key] + x + y))

Который дает следующий output.txt:

test1 ba ab cd dh gf 123 344 123
test1 ba ab cd dh gf 234 567 787
test1 ba ab cd dh gf 221 344 566
test3 rt ty er wq ee 456 121 677

Примечание: вероятно, легче следовать подходу Брэда Соломона к Пандам, поскольку это можно сделать одной командой.