如何基于密钥合并两个元组列表？

Question

我有两个需要合并的元组列表.这与数据库术语中的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))

如何合并/加入这两个列表并获得更好的性能？

Answer 1

使用@CoryKramer在评论中建议的第一个(在每个列表中唯一)列中连接两个元组列表:itertools.groupby()

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:]

这是@Padraic Cunningham提到的collections.defaultdict基于解决方案的解决方案

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]