Question

我有两个列表，我必须从第一个列表中提取项目，第二个列表中存在第一个元素。我粘贴的代码完全可以正常工作，但是当我使用数百万条记录时，它的速度非常慢。有没有人知道如何优化它？

a = [[1,0],[2,0],[3,0],[4,0]]
b = [2,4,7,8]

same_nums = list(set([x[0] for x in a]).intersection(set(b)))

result = []

for i in a:
    if i[0] in same_nums:
        result.append(i)

print(result)

Answer 1

使用过滤器尝试列表理解，而不是创建具有相同大小的第二个列表并通过集合，例如：

[x for x in a if x[0] in b]

可能会更快。

在您发布的代码中，您执行了很多操作：创建删除第二个维度的副本，从中创建一个集合，将其与另一个集合相交，然后将该集合复制回列表。这个大名单至少以这种方式处理了四次。我的建议只列出一次，所以我实际上期望它更快。

Answer 2

你太复杂了。只需将b转换为set即可加快包含检查速度。然后，理解中的a的一次迭代就足够了：

set_b = set(b)  # makes   vvvvvvvvvvvvv  O(1)
result = [x for x in a if x[0] in set_b]

特别将same_nums转回list是一个真正的性能杀手，因为它再次使整个事件O(m*n)。使用b中的单个集合，O(m+n)。但是，same_nums完全没必要，因为您知道所有i[0]都在a，因为您正在迭代a。

Answer 3

我建议使用Numpy库，因为它使用C / C ++实现，以便它运行得更快。

import numpy as np

a = np.random.randint(10000, size=(10000000, 2))
b = np.random.randint(10000, size=1000)

mask = np.isin(a[:,0], b)
a_masked = a[mask, :]

# Exec time: 3.2 s ± 185 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

相同数量元素的另一种方法（Python列表）执行时间如下：

[x for x in a if x[0] in b]

# Exec time: 55.1 s ± 5.47 s per loop (mean ± std. dev. of 7 runs, 1 loop each)

正如您所看到的，numpy的运行速度要快得多。您也可以将python列表转为numpy，反之亦然simply。

Answer 4

最优解决方案（至少我的具体情况）是 schwobaseggl 和 Novin Shahroudi 提供的解决方案

我正在使用schwobaseggl的建议，因为我正在从SQL查询中读取数据，而Novin Shahroudi解决方案要求我进行更多的数据类型转换。

Answer 5

a = [[1,0],[2,0],[3,0],[4,0]]
b = [2,4,7,8]
def return_common_elemens(a,b):
        for subval in a:
            if subval[0] in b:
                yield subval

print(list(return_common_elemens(a,b)))
>>>[[2, 0], [4, 0]]

从列表中的子列表中查找公共元素

5 个答案: