python - 从列表中的子列表中查找公共(public)元素

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

你把事情搞得太复杂了。只需将 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 是真正的性能 killer ，因为它使整个事情再次成为 O(m*n)。对于来自 b 的单个集合，它是 O(m+n)。但是 same_nums 一开始是完全没有必要的，因为您知道所有 i[0] 都在 a 中，因为您正在迭代 a 。