python - 如何有效地找到两个列表中匹配元素的索引

Question

我正在处理两个大型数据集，我的问题如下。

假设我有两个列表:

list1 = [A,B,C,D]

list2 = [B,D,A,G]

除了 O(n²) 搜索之外，如何使用 Python 高效地找到匹配的索引？结果应如下所示:

matching_index(list1,list2) -> [(0,2),(1,0),(3,1)]

Answer 1

无重复

如果您的对象是可散列的并且您的列表没有重复项，您可以创建第一个列表的倒排索引，然后遍历第二个列表。这只遍历每个列表一次，因此是 O(n)。

def find_matching_index(list1, list2):

    inverse_index = { element: index for index, element in enumerate(list1) }

    return [(index, inverse_index[element])
        for index, element in enumerate(list2) if element in inverse_index]

find_matching_index([1,2,3], [3,2,1]) # [(0, 2), (1, 1), (2, 0)]

重复

您可以扩展之前的解决方案以解决重复问题。您可以使用 set 跟踪多个索引。

def find_matching_index(list1, list2):

    # Create an inverse index which keys are now sets
    inverse_index = {}

    for index, element in enumerate(list1):

        if element not in inverse_index:
            inverse_index[element] = {index}

        else:
            inverse_index[element].add(index)

    # Traverse the second list    
    matching_index = []

    for index, element in enumerate(list2):

        # We have to create one pair by element in the set of the inverse index
        if element in inverse_index:
            matching_index.extend([(x, index) for x in inverse_index[element]])

    return matching_index

find_matching_index([1, 1, 2], [2, 2, 1]) # [(2, 0), (2, 1), (0, 2), (1, 2)]

不幸的是，这不再是O(n)。考虑输入 [1, 1] 和 [1, 1] 的情况，输出是 [(0, 0), (0, 1) , (1, 0), (1, 1)]。因此，根据输出的大小，最坏的情况不会比 O(n^2) 好。

尽管如此，如果没有重复，这个解决方案仍然是 O(n)。

不可散列的对象

现在出现了您的对象不可散列但可比较的情况。这里的想法是以保留每个元素的原始索引的方式对列表进行排序。然后我们可以对相等的元素序列进行分组以获得匹配索引。

由于我们在下面的代码中大量使用了 groupby 和 product，所以我让 find_matching_index 返回一个生成器以提高长列表的内存效率.

from itertools import groupby, product

def find_matching_index(list1, list2):
    sorted_list1 = sorted((element, index) for index, element in enumerate(list1))
    sorted_list2 = sorted((element, index) for index, element in enumerate(list2))

    list1_groups = groupby(sorted_list1, key=lambda pair: pair[0])
    list2_groups = groupby(sorted_list2, key=lambda pair: pair[0])

    for element1, group1 in list1_groups:
        try:
            element2, group2 = next(list2_groups)
            while element1 > element2:
                (element2, _), group2 = next(list2_groups)

        except StopIteration:
            break

        if element2 > element1:
            continue

        indices_product = product((i for _, i in group1), (i for _, i in group2), repeat=1)

        yield from indices_product

        # In version prior to 3.3, the above line must be
        # for x in indices_product:
        #     yield x

list1 = [[], [1, 2], []]
list2 = [[1, 2], []]

list(find_matching_index(list1, list2)) # [(0, 1), (2, 1), (1, 0)]

事实证明，时间复杂度并没有受到太大影响。排序当然需要 O(n log(n))，但是 groupby 提供的生成器可以通过仅遍历我们的列表两次来恢复所有元素。结论是我们的复杂性主要受 product 输出大小的限制。因此给出算法为 O(n log(n)) 的最佳情况和再次为 O(n^2) 的最坏情况。