设置比较优化

Question

描述

我有两个大的集合列表

A = [ {...}, ..., {...} ]
B = [ {...}, ..., {...} ]

我正在执行一个成本非常高的列表理解，对于 A 中每个集合中的每个元素，检查是否与 B 集合中的任何元素匹配，如果匹配则返回 B 各自的集合。

[find_sets(i) for i in A]

例子

一个最小的例子是这样的：

import secrets

# create sample data 
def generate_random_strings(num_strings, string_length):
    random_strings = []
    for i in range(num_strings):
        random_strings.append(secrets.token_hex(string_length))
    random_strings = set(random_strings)
    return random_strings

A = [generate_random_strings(5, 1) for i in range(10000)]
B = [generate_random_strings(5, 1) for i in range(10000)]

# set checker 
def find_sets(A):
    matching_sets = []
    for b_set in B:
        if A & b_set:
            matching_sets.append(b_set)
    return matching_sets

result = [find_set(i) for i in A]

多重处理

它在我所有的 32 个 CPU 内核上明显更快：

from tqdm.contrib.concurrent import process_map

pool = multiprocessing.Pool(processes=32)
results = process_map(find_sets, A, chunksize=100)

问题

虽然对于 A 和 B 的几千个元素，列表理解在我的机器上运行得相当快，并且多处理有助于将它扩展到大约 50.000 个元素，但对于我的实际大小的每个列表中的 500.000 个元素，它变得非常慢。

有什么方法可以通过矢量化、之前对集合进行哈希处理或使用某种优化数据类型（frozensets 没有帮助）来加速我的函数代码？

Answer 1

这在我的测试中运行速度快了一个数量级：

import collections

reverse_map = collections.defaultdict(set)
for idx, elements in enumerate(B):
    for element in elements:
        reverse_map[element].add(idx)

def find_sets(A):
    union = set()
    emptyset = set()
    for element in A:
        union |= reverse_map.get(element, emptyset)
    return [B[idx] for idx in union]

Answer 2

你可以使用理解而不是find_sets：

 result = [[b for b in B if a & b] for i in A]