python - 包含 3 个总和为 n 的正整数的数组

Question

我使用 itertools combinations_with_replacement 创建了一个生成器，它返回总和为 n 的 3 个正整数的所有组合:

def combinations(n):
    for combo in combinations_with_replacement([i for i in range(1,n+1)],3):
        if sum(combo) == n:
            yield(combo)

例如combinations(7) 返回 (1, 1, 5) (1, 2, 4) (1, 3, 3) (2, 2, 3)不幸的是，当 n 值较大时，这很快就会变得非常慢。有没有一种更有效的替代方法？我尝试过使用 for 循环，尽管每次我都会得到重复的组合。提前致谢!

Answer 1

您不必获得三个数字的所有组合。您只需得到两个数字的组合，并且您知道第三个数字必须是什么。

>>> n = 100
>>> combs_of_three = [(a,b,c) for (a,b,c) in combinations_with_replacement(range(1, n+1), 3) if a+b+c == n]
>>> combs_of_two = [(a,b,n-a-b) for (a,b) in combinations_with_replacement(range(1, n+1), 2) if n-a-b >= b]
>>> combs_of_three == combs_of_two
True

这要快得多:

>>> %timeit [(a,b,c) for (a,b,c) in combinations_with_replacement(range(1, n+1), 3) if a+b+c == n]
9.97 ms ± 97.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

>>> %timeit [(a,b,n-a-b) for (a,b) in combinations_with_replacement(range(1, n+1), 2) if n-a-b >= b]
359 µs ± 2.06 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)