Как объединить два списка кортежей на основе ключа?

У меня есть два списка кортежей, которые мне нужно объединить. Это будет сопоставимо с JOIN в терминах базы данных. Порядок кортежей в каждом списке может измениться. Порядок элементов в кортеже не изменится. Количество элементов в A должно равняться счету в B, но может быть разница.

Вот мои два списка кортежей. В каждом списке будет 10 000 + этих кортежей, поэтому производительность вызывает беспокойство. Первый элемент в каждом кортеже является ключевым для каждого списка.

listA = [(u'123', u'a1', u'a2', 123, 789), (u'124', u'b1', u'b2', 456, 357), (u'125', u'c1', u'c2', 156, 852)] listB = [(u'125', u'd1', u'N', u'd2', 1), (u'123', u'f1', u'Y', u'f2', 2)] 

Желаемый результат:

 listC = [(u'123', u'a1', u'a2', 123, 789, u'f1', u'Y', u'f2', 2), (u'125', u'c1', u'c2', 156, 852, u'd1', u'N', u'd2', 1)] 

Вот код, который я собрал вместе для тестирования концепции. Он работает, но, как вы видите, производительность – это проблема. Производительность этого кода при работе с реальными данными (10 тыс. Элементов в каждом списке) неприемлема, так как потребуется, возможно, часов для завершения.

Вот код:

 for row in listA: for item in listB: if item[0] == row[0]: item = list(item) del item[0] row = list(row) merged.append(tuple(row + item)) 

Как я могу объединить / присоединиться к двум спискам и добиться лучшей производительности?

Внутреннее объединение двух списков кортежей на первом (уникальном в каждом списке) столбце с использованием itertools.groupby() предложенного @CoryKramer в комментариях :

 from itertools import groupby from operator import itemgetter def inner_join(a, b): L = a + b L.sort(key=itemgetter(0)) # sort by the first column for _, group in groupby(L, itemgetter(0)): row_a, row_b = next(group), next(group, None) if row_b is not None: # join yield row_a + row_b[1:] # cut 1st column from 2nd row 

Пример:

 result = list(inner_join(listA, listB)) assert result == listC 

Это решение имеет сложность времени O(n*log n) (ваше решение (в вопросе) – O(n*n) что намного хуже при n ~ 10000 ).

Это не имеет значения для небольшого n такого как 10**4 в вопросе, но в Python 3.5+ вы можете использовать heapq.merge() с key параметром, чтобы избежать выделения нового списка, то есть для решения O(1) постоянной памяти:

 from heapq import merge # merge has key parameter in Python 3.5 def inner_join(a, b): key = itemgetter(0) a.sort(key=key) b.sort(key=key) for _, group in groupby(merge(a, b, key=key), key): row_a, row_b = next(group), next(group, None) if row_b is not None: # join yield row_a + row_b[1:] # cut 1st column from 2nd row 

Вот диктофонное решение. O(n) линейный по времени и пространству алгоритм:

 def inner_join(a, b): d = {} for row in b: d[row[0]] = row for row_a in a: row_b = d.get(row_a[0]) if row_b is not None: # join yield row_a + row_b[1:] 

Вот решение, основанное на collections.defaultdict упомянутое @Padraic Cunningham

 from collections import defaultdict from itertools import chain def inner_join(a, b): d = defaultdict(list) for row in chain(a, b): d[row[0]].append(row[1:]) for id, rows in d.iteritems(): if len(rows) > 1: assert len(rows) == 2 yield (id,) + rows[0] + rows[1] 

Вы раньше использовали панды? Это, похоже, дает желаемый результат:

 n [41]: import pandas as pd listA = [(u'123', u'a1', u'a2', 123, 789), (u'124', u'b1', u'b2', 456, 357), (u'125', u'c1', u'c2', 156, 852)] listB = [(u'125', u'd1', u'N', u'd2', 1), (u'123', u'f1', u'Y', u'f2', 2)] A = pd.DataFrame(listA) B = pd.DataFrame(listB) A.merge(B, on=0) Out[41]: 0 1_x 2_x 3_x 4_x 1_y 2_y 3_y 4_y 0 123 a1 a2 123 789 f1 Y f2 2 1 125 c1 c2 156 852 d1 N d2 1 

«A» и «B» – это фреймы данных pandas, в которые встроена часть встроенных в SQL функций, таких как слияние. Если вы не использовали панды, дайте мне знать, если вам нужны дополнительные объяснения.

См. Присоединение / слияние DataFrame в стиле базы данных .

Вы можете группировать первый элемент с помощью OrderedDict, добавляя каждый кортеж, а затем сохраняйте и присоединяйте кортежи, где список значений имеет длину> 1:

 from itertools import chain from collections import OrderedDict od = OrderedDict() for ele in chain(listA,listB): od.setdefault(ele[0], []).append(ele[1:]) print([(k,) + tuple(chain.from_iterable(v)) for k,v in od.iteritems() if len(v) > 1]) 

Вывод:

 [('123', 'a1', 'a2', 123, 789, 'f1', 'Y', 'f2', 2), ('125', 'c1', 'c2', 156, 852, 'd1', 'N', 'd2', 1)] 

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

Или сохраните объекты itertools.islice используя флаг для поиска совпадающих ключей:

 from itertools import chain, islice from collections import OrderedDict od = OrderedDict() for ele in chain(listA, listB): k = ele[0] if k in od: od[k]["v"].append(islice(ele, 1, None)) od[k]["flag"] = True else: od.setdefault(k, {"flag": False, "v": []})["v"].append(islice(ele, 1, None)) print([(k,) + tuple(chain.from_iterable(v["v"])) for k, v in od.items() if v["flag"]]) 

Вывод:

 [('123', 'a1', 'a2', 123, 789, 'f1', 'Y', 'f2', 2), ('125', 'c1', 'c2', 156, 852, 'd1', 'N', 'd2', 1)]