python - 检查值或值列表是否是 python 中列表的子集的最快方法

Question

我有一个非常大的列表列表，称为 main_list，包含大约 1300 万个列表，每个列表包含 6 个数字。我正在寻找一种方法来过滤掉任何不包含特定值的列表。例如，要创建一个新的列表列表，其中仅包含值为 4 和 5 的列表，我的代码如下工作:

and_include = []
temp_list=[4,5]
for sett in main_list:
    if set(temp_list).issubset(sett):
        and_include.append(sett)

这大约需要 5 秒才能运行，这对于频繁使用来说会很烦人，所以我想知道是否有更快的方法来执行此操作，使用 numpy 或 cython？

我对 cython 不是很熟悉，但我尝试以这种方式实现、编译它，但我遇到了一个错误。

def andinclude(list main_list,list temp_list):
    and_include=[]
    for sett in main_list:
        if set(temp_list).issubset(sett):
            and_include.append(sett)
    return and_include

希望有更快的方法吗？

Answer 1

这是一个 numpy 解决方案:

import numpy as np

# Randomly generate 2d array of integers
np.random.seed(1)
a = np.random.randint(low=0, high=9, size=(13000000, 6))

# Use numpy indexing to filter rows
results = a[(a == 4).any(axis=1) & (a == 5).any(axis=1)]

结果:

In [35]: print(results.shape)
(3053198, 6)

In [36]: print(results[:5])
[[5 5 4 5 5 1]
 [5 5 4 3 8 6]
 [2 5 8 1 1 4]
 [0 5 4 1 1 5]
 [3 2 5 2 4 6]]

时间:

In [37]: %timeit results = a[(a == 4).any(axis=1) & (a == 5).any(axis=1)]
923 ms ± 38.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

如果您需要将结果转换回列表列表而不是二维 numpy 数组，您可以使用:

l = results.tolist()

这使在我的机器上运行的时间增加了大约 50%，但仍然比任何涉及循环 Python 列表的解决方案都要快。