在列表中找到子列表中的共同元素

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)

- Finrod Felagund

Answer 1

尝试使用过滤器的列表推导式，而不是创建一个与大小相同的第二个列表并通过集合进行操作，类似于以下内容：

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

可能会更快一些。

在您发布的代码中，您做了很多事情：您创建了一个去掉第二维度的副本，从中创建了一个集合，将其与另一个集合相交，并将该集合复制回一个列表。这样至少处理了四次大型列表。我的建议只需要通过列表一次，所以我实际上希望它更快。

- Demosthenes

---

惊人。我会考虑一下，也许我会想到另一个选项。 - Demosthenes

在这里，“in b” 部分是比较慢的，因为对于每个项目“x”，都需要遍历一次“b”。定义“_b = set(b)” 并在推导式中替换它。这样应该会大大加快速度。 - GPhilo

“if in” 部分是减慢程序的原因，因为我处理的数据集非常庞大。我想知道 “intersection” 方法是否可以仅检查子列表中的第一项。 - Finrod Felagund

不错的发现。我喜欢它。 - Demosthenes

@FinrodFElagund的集合是基于哈希的，但是您在代码中使用它们的方式会让您多次遍历 a（一次用于创建 set(a)，另一次用于交集）。如果您像 @schwobaseggl 在他的答案中实现的那样使用集合推导式（这就是我在上面的评论中建议的），则不需要这样做。 - GPhilo

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

- user2390182

---

是的，那快了好几千倍！非常感谢！ - Finrod Felagund

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。

- Novin Shahroudi

Answer 4

对于我的特定情况，最优解决方案是由schwobaseggl和Novin Shahroudi提供的。

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

- Finrod Felagund

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

- Veera Balla Deva

---

感谢您提供这段代码片段，它可能提供了一些有限的短期帮助。通过展示为什么这是一个好的问题解决方案，适当的解释将极大地改善其长期价值，并使其对未来读者的其他类似问题更有用。请编辑您的答案以添加一些解释，包括您所做的假设。 - Toby Speight