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

user_5218354 0
Viewed 568

How would I efficiently find all the possible combinations of n positive integers adding up to a given number k in Python?

I know I could solve this by filtering all the possible combinations:

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)]

I have seen something similar has been discussed in abstract way here, but I do not see a simple way of translating this into code.

Also, I would prefer an iterative solution over a recursive one.

combinatorics algorithm python
3 Answers

A much more efficient than OP, but still filter-based (and hence less efficient than the accepted answer) approach is to use:

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)]

Some sample timings:

%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

(Disclaimer: I am the main author of the flyingcircus package).

user_52183540

A recursive solution based on 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)]

And an iterative one:

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)]

Time-wise, the recursive solution seems to be fastest:

%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
user_52183540

Here's a recursive solution: 

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)

results in: 

[[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],

And the same solution can be optimized to a semi-dynamic programming solution by keeping a 'cache' of all the results that we're previously computed, and reusing them when needed: 

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
user_11055600
Related