python - 如何提高涉及生成器和嵌套 for 循环的代码的效率

Question

我正在编写一个代码，该代码基于两个均按升序排列的数字列表(a 和 b)以渐进方式生成子列表列表。每个包含两个元素的子列表可以被视为这两个列表中的元素的组合。第二个元素(来自列表 b)需要大于第一个元素(来自列表 a)。特别地，对于第二个元素，该值可能并不总是数字。子列表可以是 [elem, None]，这意味着列表 b 中没有与列表 a 中的“elem”匹配的内容。最终输出中不应有任何重复项。如果您想象输出位于表格中，则每个子列表将是一行，并且在两列中的每一列中，除了第二列中的“无”之外，元素均按升序排列。

由于我上一个问题的友好回复，我受到了启发并编写了可以实现目标的代码。 ( How to generate combinations with none values in a progressive manner ) 代码如下所示。

import itertools as it
import time

start=time.time()

a=[1,5,6,7,8,10,11,13,15,16,20,24,25,27]
b=[2,8,9,10,11,12,13,14,17,18,21,26]

def create_combos(lst1, lst2): #a is the base list; l is the adjacent detector list
    n = len(lst1)
    x_ref = [None,None]
    for i in range(1,n+1):
        choices_index = it.combinations(range(n),i)
        choices_value = list(it.combinations(lst2,i)) 
        for choice in choices_index:
            for values in choices_value:
                x = [[elem,None] for elem in lst1]
                for index,value in zip(choice,values): #Iterate over two lists in parallel  
                    if value <= x[index][0]:
                        x[index][0] = None
                        break
                    else:
                        x[index][1] = value #over-write in appropriate location
                if x_ref not in x:
                    yield x

count=0
combos=create_combos(a,b)
for combo in combos:
#    print(combo)
    count+=1
print('The number of combos is ',count)

end=time.time()
print('Run time is ',end-start)

这段代码是我以有限的 Python 知识在速度方面能得到的最好的代码。然而，随着列表 a 和 b 中的元素数量超过 15，运行时间仍然太长。我理解这可能是因为组合的急剧增加。然而，我想知道是否可以进行任何改进来提高其效率，也许是关于生成组合的方式。此外，我生成了所有可能的组合，然后删除了不合适的组合，我认为这也可能效率低下。

期望的结果是在合理的时间内处理每个列表中的大约 30 个元素。

编辑:由于一旦每个列表中的元素数量增加，组合的数量也会急剧增加。因此，我想保留生成器，以便一次只允许一个组合占用内存。

如果我对上述任何陈述有不清楚的地方，请随时提问。谢谢:)

Answer 1

编辑2:

好吧，如果你做得更聪明一点，你就能更快地完成这件事。我现在将使用 NumPy 和 Numba 来真正加速速度。如果您不想使用 Numba，只需注释使用它的部分，它仍然可以工作，只是速度较慢。如果您不需要 NumPy，可以将其替换为列表或嵌套列表，但同样可能会带来显着的性能差异。

所以让我们看看。需要改变的两个关键事项是:

• 为输出预先分配空间(我们不使用生成器，而是立即生成整个输出)。
• 重复使用计算出的组合。

要进行预分配，我们需要首先计算总共有多少种组合。该算法类似，但只是计数，如果您有部分计数的缓存，那么实际上速度相当快。 Numba 确实在这里发挥了巨大的作用，但我已经使用过它。

import numba as nb

def count_combos(a, b):
    cache = np.zeros([len(a), len(b)], dtype=np.int32)
    total = count_combos_rec(a, b, 0, 0, cache)
    return total

@nb.njit
def count_combos_rec(a, b, i, j, cache):
    if i >= len(a) or j >= len(b):
        return 1
    if cache[i][j] > 0:
        return cache[i][j]
    while j < len(b) and a[i] >= b[j]:
        j += 1
    count = 0
    for j2 in range(j, len(b)):
        count += count_combos_rec(a, b, i + 1, j2 + 1, cache)
    count += count_combos_rec(a, b, i + 1, j, cache)
    cache[i][j] = count
    return count

现在我们可以为所有组合预先分配一个大数组。我不会直接将组合存储在那里，而是拥有一个整数数组，表示 b 中元素的位置(a 中的元素由位置隐式表示，并且None 匹配由 -1 表示)。

为了重用组合，我们执行以下操作。每次我们需要找到某个对i/j的组合时，如果之前没有计算过，我们就这样做，然后将位置保存在第一次存储这些组合的输出数组。下次遇到相同的i/j对时，我们只需复制之前制作的相应部分即可。

总而言之，算法最终如下(本例中的结果是一个 NumPy 对象数组，第一列是来自 a 的元素，第二列是来自 b 的元素，但您可以使用 .tolist() 将其转换为常规 Python 列表)。

import numpy as np
import numba as nb

def generate_combos(a, b):
    a = np.asarray(a)
    b = np.asarray(b)
    # Count combos
    total = count_combos(a, b)
    count_table = np.zeros([len(a), len(b)], np.int32)
    # Table telling first position of a i/j match
    ref_table = -np.ones([len(a), len(b)], dtype=np.int32)
    # Preallocate result
    result_idx = np.empty([total, len(a)], dtype=np.int32)
    # Make combos
    generate_combos_rec(a, b, 0, 0, result_idx, 0, count_table, ref_table)
    # Produce matchings array
    seconds = np.where(result_idx >= 0, b[result_idx], None)
    firsts = np.tile(a[np.newaxis], [len(seconds), 1])
    return np.stack([firsts, seconds], axis=-1)

@nb.njit
def generate_combos_rec(a, b, i, j, result, idx, count_table, ref_table):
    if i >= len(a):
        return idx + 1
    if j >= len(b):
        result[idx, i:] = -1
        return idx + 1
    elif ref_table[i, j] >= 0:
        r = ref_table[i, j]
        c = count_table[i, j]
        result[idx:idx + c, i:] = result[r:r + c, i:]
        return idx + c
    else:
        idx_ref = idx
        j_ref = j
        while j < len(b) and a[i] >= b[j]:
            j += 1
        for j2 in range(j, len(b)):
            idx_next = generate_combos_rec(a, b, i + 1, j2 + 1, result, idx, count_table, ref_table)
            result[idx:idx_next, i] = j2
            idx = idx_next
        idx_next = generate_combos_rec(a, b, i + 1, j, result, idx, count_table, ref_table)
        result[idx:idx_next, i] = -1
        idx = idx_next
        ref_table[i, j_ref] = idx_ref
        count_table[i, j_ref] = idx - idx_ref
        return idx

让我们检查一下结果是否仍然正确:

a = [1, 5, 6, 7, 8, 10, 11, 13, 15, 16, 20, 24, 25, 27]
b = [2, 8, 9, 10, 11, 12, 13, 14, 17, 18, 21, 26]
# generate_combos_prev is the previous recursive method
combos1 = list(generate_combos_prev(a, b))
# Note we do not need list(...) here because it is not a generator
combos = generate_combos(a, b)
print((combos1 == combos).all())
# True

好的，现在让我们看看性能。

%timeit list(generate_combos_prev(a, b))
# 3.7 s ± 17.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit generate_combos(a, b)
# 593 ms ± 2.66 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

不错!速度快了 6 倍!除了额外的依赖性之外，唯一可能的缺点是我们一次而不是迭代地进行所有组合(因此您将同时将它们全部存储在内存中)并且我们需要一个大小为 O(len(a) * len(b))。

<小时/>

这是一种更快的方式来完成您正在做的事情:

def generate_combos(a, b):
    # Assumes a and b are already sorted
    yield from generate_combos_rec(a, b, 0, 0, [])

def generate_combos_rec(a, b, i, j, current):
    # i and j are the current indices for a and b respectively
    # current is the current combo
    if i >= len(a):
        # Here a copy of current combo is yielded
        # If you are going to use only one combo at a time you may skip the copy
        yield list(current)
    else:
        # Advance j until we get to a value bigger than a[i]
        while j < len(b) and a[i] >= b[j]:
            j += 1
        # Match a[i] with every possible value from b
        for j2 in range(j, len(b)):
            current.append((a[i], b[j2]))
            yield from generate_combos_rec(a, b, i + 1, j2 + 1, current)
            current.pop()
        # Match a[i] with None
        current.append((a[i], None))
        yield from generate_combos_rec(a, b, i + 1, j, current)
        current.pop()

a = [1, 5, 6, 7, 8, 10, 11, 13, 15, 16, 20, 24, 25, 27]
b = [2, 8, 9, 10, 11, 12, 13, 14, 17, 18, 21, 26]
count = 0
combos = generate_combos(a, b)
for combo in combos:
    count += 1
print('The number of combos is', count)
# 1262170

此算法的唯一区别在于，它比您的多生成一个组合(在您的代码中，最终计数为 1262169)，即 a 中的每个元素都与 None 匹配的组合。这始终是最后生成的组合，因此您可以根据需要忽略该组合。

编辑:如果您愿意，可以将 generate_combos_rec 中的 # Match a[i] with None block 移动到 while 之间循环和 for 循环，然后 a 中每个值与 None 匹配的额外组合将是第一个而不是最后一个。这可能会更容易跳过它。或者，您可以将 yield list(current) 替换为:

if any(m is not None for _, m in current):
    yield list(current)

避免生成额外的组合(以对每个生成的组合进行额外检查为代价)。

编辑2:

这是一个稍微修改过的版本，它通过在递归中携带 bool 指示符来避免额外的组合。

def generate_combos(a, b):
    yield from generate_combos_rec(a, b, 0, 0, [], True)

def generate_combos_rec(a, b, i, j, current, all_none):
    if i >= len(a):
        if not all_none:
            yield list(current)
    else:
        while j < len(b) and a[i] >= b[j]:
            j += 1
        for j2 in range(j, len(b)):
            current.append((a[i], b[j2]))
            yield from generate_combos_rec(a, b, i + 1, j2 + 1, current, False)
            current.pop()
        current.append((a[i], None))
        yield from generate_combos_rec(a, b, i + 1, j, current, all_none)
        current.pop()