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

Question

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.

what kind of answer are you expecting? what's wrong with your code? — Nicolas Gervais, Jun 12 '20 at 12:53
You could make that code rather more efficient by only finding the combinations of n-1 numbers - and then *calculating* what the final number would have to be to make the desired total (rejecting any combination where that number would be non-positive). — jasonharper, Jun 12 '20 at 12:54
@NicolasGervais It is utterly inefficien. Only a minor portion of all the generated tuples satisfy the condition, and this gets worse with larger inputs. `list(to_sum_k(4, 100))` takes ~15 sec to produce just ~150k entries, out of ~90m seen. — norok2, Jun 12 '20 at 14:32
you're right. in this case you might be interested in know that your original solution took 23.5 seconds on my computer and my solution took 1.43 seconds. your answer to your own post with your own package yielded a time of 1.51 seconds — Nicolas Gervais, Jun 12 '20 at 14:48
@NicolasGervais except that your solution is missing quite many results, that mine isn't. — norok2, Jun 12 '20 at 14:59

norok2 · Answer 1 · 2020-10-05T22:11:29.107

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

I just added this to avoid somebody else suggesting this, because this is not what I want. — norok2, Jun 13 '20 at 08:08

norok2 · Accepted Answer · 2020-10-04T21:00:47.560

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

score 1 · Answer 3 · answered Jun 12 '20 at 12:58

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

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

3 Answers3

Linked