代码之家 › 专栏 › 技术社区 › norok2

在Python中找到“n”个正数加起来等于“k”的所有组合?

combinatorics algorithm python

norok2 · 技术社区 · 4 年前

我如何有效地找到所有可能的组合 n 与给定数相加的正整数 k 在Python中?

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

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 ,但我没有看到将其转换为代码的简单方法。

此外,我更喜欢迭代解而不是递归解。

0 回复 | 直到 4 年前

norok2 4 年前

基于的递归解 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

norok2 4 年前

一种比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 包装)。

Roy2012 4 年前

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 = {}
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