如何高效地在两个列表中查找匹配元素的索引

Question

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

假设我有两个列表：

list1 = [A,B,C,D]

list2 = [B,D,A,G]

如何有效地使用Python查找匹配的索引，而不是O(n ²)的搜索？结果应如下所示：

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

- Haoran

---

这些元素是否可哈希，例如字符串，可以用作dict中的键？ - tdelaney

是的，这些元素是可哈希的。现在我有的是：[i for i, item in enumerate(list(df1)) if item in set(list(df2))]。这会给出df1匹配索引的列表，但我想知道确切的匹配位置在哪里。非常感谢！ - Haoran

我在这里没有一个具体的答案，但是numpy在这种情况下非常出色。 - Jim Stewart

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)。

- Olivier Melançon

---

非常感谢您的帮助。是的，这正是我一直在苦苦挣扎的问题。 - Haoran

有没有办法解决重复值的问题？例如：list1 = [A,B,C,D,E] list2 = [B,A,D,A,G] ->[(0,1),(0,3),(1,0),(3,2)]? - Haoran

Answer 2

如果您的对象不可哈希，但仍可排序，则可以考虑使用 sorted 来匹配两个列表。

假设两个列表中的所有元素均有匹配项

您可以对列表索引进行排序并将结果配对。

indexes1 = sorted(range(len(list1)), key=lambda x: list1[x])
indexes2 = sorted(range(len(list2)), key=lambda x: list2[x])
matches = zip(indexes1, indexes2)

如果不是所有元素都匹配，但每个列表内部没有重复项

您可以同时对两个列表进行排序，并在排序时保留索引。然后，如果捕获到任何连续的重复项，则知道它们来自不同的列表。

biglist = list(enumerate(list1)) + list(enumerate(list2))
biglist.sort(key=lambda x: x[1])
matches = [(biglist[i][0], biglist[i + 1][0]) for i in range(len(biglist) - 1) if biglist[i][1] == biglist[i + 1][1]]

- Fred

---

好的，虽然这是O(n log(n))。 - Olivier Melançon

我看到大家都在用字典，所以想带来一些不同的东西。 :) - Fred

当然，如果对象不可哈希，则这可能很有用！你应该提到这一点。 - Olivier Melançon

这实际上是一个我没有想到的非常好的理由，哈哈。 - Fred

非常感谢。确实，这对于不可哈希的对象非常好。感谢您的想法！ - Haoran

Answer 3

针对这个问题的一种粗暴答案（即使没有其他原因，也是为了验证任何解决方案）如下：

[(xi, xp) for (xi, x) in enumerate(list1) for (xp, y) in enumerate(list2) if x==y]

您需要优化的方式在很大程度上取决于数据量和内存容量，因此了解这些列表的规模可能会有所帮助。我想下面讨论的方法至少适用于具有数百万个值的列表。

由于字典访问是O(1)的，似乎值得尝试将第二个列表中的元素映射到它们的位置。假设可以重复相同的元素，则使用collections.defaultdict可以轻松地构建必要的字典。

l2_pos = defaultdict(list)
for (p, k) in enumerate(list2):
    l2_pos[k].append(p)

表达式l2_pos[k]现在是一个列表，其中包含元素k在list2中出现的位置。现在只需要将这些位置与list1中相应键的位置配对即可。最终的结果以列表形式呈现。

[(p1, p2) for (p1, k) in enumerate(list1) for p2 in l2_pos[k]]

如果这些结构较大，你最好使用生成器表达式。要将名称绑定到上述列表推导式中的表达式，则应编写：

values = ((p1, p2) for (p1, k) in enumerate(list1) for p2 in l2_pos[k])

如果您遍历values，则可以避免创建包含所有值的列表所带来的开销，从而减少对Python内存管理和垃圾回收的负载，这几乎是与解决问题相关的所有开销。

当你开始处理大量数据时，了解生成器可能意味着有足够的内存来解决问题或者没有足够的内存。在许多情况下，它们比列表理解更具明显的优势。

编辑：除非排序变化会有害，否则使用集合而不是列表来保存位置可以进一步加速此技术。该更改留给读者自己练习。

- holdenweb

Answer 4

使用 dict 可以减少查找时间，而 collections.defaultdict 的专业化可以帮助进行簿记。目标是一个 dict，其值是您要查找的索引对。重复的值会覆盖列表中较早的值。

import collections

# make a test list
list1 = list('ABCDEFGHIJKLMNOP')
list2 = list1[len(list1)//2:] + list1[:len(list1)//2]

# Map list items to positions as in: [list1_index, list2_index]
# by creating a defaultdict that fills in items not in list1,
# then adding list1 items and updating with with list2 items. 
list_indexer = collections.defaultdict(lambda: [None, None],
 ((item, [i, None]) for i, item in enumerate(list1)))
for i, val in enumerate(list2):
    list_indexer[val][1] = i

print(list(list_indexer.values()))

- tdelaney

Answer 5

这里有一个使用 defaultdict 的简单方法。

假设

import collections as ct


lst1 = list("ABCD")
lst2 = list("BDAG")
lst3 = list("EAB")
str1 = "ABCD"

代码

def find_matching_indices(*iterables, pred=None):
    """Return a list of matched indices across `m` iterables."""
    if pred is None:
        pred = lambda x: x[0]

    # Dict insertion
    dd = ct.defaultdict(list)
    for lst in iterables:                                          # O(m)
        for i, x in enumerate(lst):                                # O(n)
            dd[x].append(i)                                        # O(1)

    # Filter + sort
    vals = (x for x in dd.values() if len(x) > 1)                  # O(n)
    return sorted(vals, key=pred)                                  # O(n log n)

演示

在两个列表中查找匹配项（根据原始帖子）：

find_matching_indices(lst1, lst2)
# [[0, 2], [1, 0], [3, 1]]

按不同的结果索引排序：

find_matching_indices(lst1, lst2, pred=lambda x: x[1])
# [[1, 0], [3, 1], [0, 2]]

匹配两个或多个迭代对象中的项目（可选择具有可变长度）：

find_matching_indices(lst1, lst2, lst3, str1)
# [[0, 2, 1, 0], [1, 0, 2, 1], [2, 2], [3, 1, 3]]

---

细节

字典插入

每个项目都附加到defaultdict的列表中。结果看起来像这样，稍后会进行过滤：

defaultdict(list, {'A': [0, 2], 'B': [1, 0], 'C': [2], 'D': [3, 1], 'G': [3]})

乍一看，从双重for循环中，我们可能会认为时间复杂度是O(n²)。然而，外部循环中的容器列表具有长度m。内部循环处理每个长度为n的容器的元素。我不确定最终的复杂度是什么，但基于this answer，我怀疑它是O(n*m)，或至少低于O(n²)。

过滤

过滤掉非匹配项（长度为1的列表），并对结果进行排序（主要用于Python<3.6中无序字典）。

使用timsort算法通过sorted按某个索引对字典值（列表）进行排序，最坏情况是O(n log n)。由于Python 3.6+中保留了字典键插入，预排序的项将减少复杂度O(n)。

总体而言，如果在Python < 3.6中使用sorted，则最佳情况时间复杂度为O(n)，最坏情况为O(n log n)，否则为O(n*m)。

- pylang