使用 itertools.product 时如何跳过一些迭代？答案

【问题标题】：How to skip some iterations when using itertools.product?使用 itertools.product 时如何跳过一些迭代？
【发布时间】：2021-11-02 04:06:16
【问题描述】：

假设有三个排序列表，A、B、C。

A = [1, 2, 3, 4]
B = [3, 4, 5]
C = [2, 3, 4]

我正在使用itertools.product 查找总和小于 10 的所有可能组合。

如果我只有三个列表，我将使用以下代码。

A = [1, 2, 3, 4] B = [3, 4, 5] C = [2, 3, 4]

for a in A:
    for b in B:
        for c in C:
            if a + b + c < 10:
                print(a, b, c)
            else:
                break

这里，每个列表都是排序的，因此我使用 break 来提高效率。

但是当我使用itertools.product 时，那我如何使用break 呢？我的意思是如何直接进行特定的迭代（例如，a = 3，b = 3，c = 3）？

for a, b, c in itertools.product(A, B, C):
   ....?

【问题讨论】：

使用continue 跳到下一个迭代。
@ddejohn 我认为这不是 OP 想要的。他/她想跳过多次迭代。
啊，我现在明白了。 OP，使用itertools.product 做你想做的事是不可能的。

标签： python for-loop iteration itertools

【解决方案1】：

您可以尝试以下方法：

from itertools import product, dropwhile

A = [1, 2, 3, 4] 
B = [3, 4, 5] 
C = [2, 3, 4]

for a, b, c in dropwhile(lambda x: x != (3,3,3), product(A, B, C)):
    print(a, b, c)

它给出：

请注意，这并没有真正直接进入给定的迭代。相反，itertools.dropwhile 运行迭代器，直到满足指定条件，然后才开始返回其值。

【讨论】：

【解决方案2】：

不可能在itertools.product 中跳过迭代，但鉴于列表已排序，可以通过使用二分搜索并查找低于所需差异的项目并使用记忆化来减少迭代次数：

import itertools
import bisect


def bisect_fast(A, B, C, threshold):
    seen_b_diff = {}
    seen_c_diff = {}

    for a in A:
        b_diff = threshold - a
        if b_diff not in seen_b_diff:
            index =  bisect.bisect_left(B, b_diff)
            seen_b_diff[b_diff] = index

        # In B we are only interested in items that are less than `b_diff`
        for ind in range(seen_b_diff[b_diff]):
            b = B[ind]
            c_diff = threshold - (a + b)
            # In `C` we are only interested in items that are less than `c_diff`
            if c_diff not in seen_c_diff:
                index = bisect.bisect_left(C, c_diff)
                seen_c_diff[c_diff] = index

            for ind in range(seen_c_diff[c_diff]):
                yield a, b, C[ind] 


def naive(A, B, C, threshold):
    for a, b, c in itertools.product(A, B, C):
        if a + b + c < threshold:
            yield a, b, c

输出

>>> from random import choice
>>> A, B, C = [sorted([choice(list(range(1000))) for _ in range(250)]) for _ in range(3)]
>>> list(naive(A, B, C, 1675)) == list(bisect_fast(A, B, C, 1675))
True
>>> %timeit list(bisect_fast(A, B, C, 1675))
1.59 s ± 32.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
>>> %timeit list(naive(A, B, C, 1675))
3.09 s ± 55.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

【讨论】：