在 Python 中找出加起来等于 `k` 的 `n` 个正数的所有组合？

Question

我如何有效地找到 Python 中 n 个正整数加起来等于给定数字 k 的所有可能组合？

我知道我可以通过过滤所有可能的组合来解决这个问题：

import itertools


def to_sum_k(n, k):
    for i in itertools.product(range(1, k - n + 2), repeat=n):
        if sum(i) == k:
            yield i


print(list(to_sum_k(3, 5)))
# [(1, 1, 3), (1, 2, 2), (1, 3, 1), (2, 1, 2), (2, 2, 1), (3, 1, 1)]

我看到有人以抽象的方式讨论了类似的内容 here，但我没有看到将其转化为代码的简单方法。

此外，与递归解决方案相比，我更喜欢迭代解决方案。

Answer 1

这是一个递归的解决方案：

def to_sum_k(n, k):
    if n == 1: 
        return [ [k] ]
    if n > k or n <= 0:
        return []
    res = []
    for i in range(k):
        sub_results = to_sum_k(n-1, k-i)
        for sub in sub_results:
            res.append(sub + [i])
    return res    

to_sum_k(3, 5)

结果：

[[5, 0, 0],
 [4, 1, 0],
 [3, 2, 0],
 [2, 3, 0],
 [1, 4, 0],
 [4, 0, 1],
 [3, 1, 1],
 [2, 2, 1],
 [1, 3, 1],
 [3, 0, 2],
 ...
 [2, 1, 2],

通过保留我们之前计算的所有结果的 'cache' 并在需要时重复使用它们，可以将相同的解决方案优化为半动态规划解决方案：

cache = {}
def to_sum_k(n, k):
    res = cache.get((n,k), [])
    if res: 
        return res

    if n == 1: 
        res  = [ [k] ]
    elif n > k or n <= 0:
        res = []
    else:
        for i in range(k):
            sub_results = to_sum_k(n-1, k-i)
            for sub in sub_results:
                res.append(sub + [i])
    cache[(n,k)] = res
    return res

Answer 2

一种比 OP 更有效，但仍然基于过滤器（因此效率低于公认的答案）的方法是使用：

import itertools
import flyingcircus as fc


def to_sum_k(n, k):
    for i in itertools.combinations_with_replacement(range(1, k - n + 2), r=n):
        if sum(i) == k:
            yield from fc.unique_permutations(i)


print(list(to_sum_k(3, 5)))
# [(1, 1, 3), (1, 2, 2), (1, 3, 1), (2, 1, 2), (2, 2, 1), (3, 1, 1)]

一些示例时间：

%timeit list(to_sum_k_OP(4, 80))
# 1 loop, best of 3: 5.43 s per loop
%timeit list(to_sum_k(4, 80))
# 1 loop, best of 3: 331 ms per loop

(免责声明：我是flyingcircus包的主要作者。

Answer 3

基于this的递归解：

def to_sum_k_rec(n, k):
    if n == 1:
        yield (k,)
    else:
        for x in range(1, k):
            for i in to_sum_k_rec(n - 1, k - x):
                yield (x,) + i


print(list(to_sum_k_rec(3, 5)))
# [(1, 1, 3), (1, 2, 2), (1, 3, 1), (2, 1, 2), (2, 2, 1), (3, 1, 1)]

还有一个迭代：

import itertools


def to_sum_k_iter(n, k):
    index = [0] * (n + 1)
    index[-1] = k
    for j in itertools.combinations(range(1, k), n - 1):
        index[1:-1] = j
        yield tuple(index[i + 1] - index[i] for i in range(n))


print(list(to_sum_k_iter(3, 5)))
# [(1, 1, 3), (1, 2, 2), (1, 3, 1), (2, 1, 2), (2, 2, 1), (3, 1, 1)]

就时间而言，递归解决方案似乎是最快的：

%timeit list(to_sum_k_OP(4, 100))
# 1 loop, best of 3: 13.9 s per loop
%timeit list(to_sum_k_rec(4, 100))
# 10 loops, best of 3: 101 ms per loop
%timeit list(to_sum_k_iter(4, 100))
# 1 loop, best of 3: 201 ms per loop

在 Python 中找出加起来等于 `k` 的 `n` 个正数的所有组合？

Find all combinations of `n` positive numbers adding up to `k` in Python?

python

algorithm

combinatorics