CODEWITHSUNDEEP
python
Jacob
Jacob

Reputation: 810

Arrays with 3 positive integers that sum to n

I have created a generator using itertools combinations_with_replacement which returns all the combinations of 3 positive integers that sum to n:

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

e.g. combinations(7) returns (1, 1, 5) (1, 2, 4) (1, 3, 3) (2, 2, 3) Unfortunately this quickly becomes very slow with larger values of n. Is there an alternative way of doing this which is more efficient? I have tried using for loops although every time I get duplicate combinations. Thanks in advance!

Upvotes: 2

Views: 133

Answers (4)

Prune
Prune

Reputation: 77847

You can do this with a recursive function: choose a number that fits; then recur on the remaining size and total.

import math

def partition(total, size=3, lowest=1):
    if size == 1:
        return [[total]]
    else:
        result = []

        # At each choice, pick no less than a "fair" proportion of the remaining total.
        #    This avoids duplicating combinations.
        limit = math.ceil(total / size)

        # Iterate `i` through the range [ limit, total-(size-1) ], inclusive
        for i in range(limit, total-size+2):
            for shorter in partition(total-i, size-1):
                result.append(shorter + [i])
    return result

print(partition( 7, 3))
print(partition(12, 3))

Output:

[[2, 2, 3], [1, 3, 3], [1, 2, 4], [1, 1, 5]]
[[4, 4, 4], [3, 5, 4], [2, 6, 4], [1, 7, 4], [3, 4, 5], [2, 5, 5], [1, 6, 5], [3, 3, 6], [2, 4, 6], [1, 5, 6], [2, 3, 7], [1, 4, 7], [2, 2, 8], [1, 3, 8], [1, 2, 9], [1, 1, 10]]

Upvotes: 1

RomanPerekhrest
RomanPerekhrest

Reputation: 92854

Another optimization with itertools.filterfalse:

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

In [384]: def combinations_2(n):
     ...:     return itertools.filterfalse(lambda x: x[0]+x[1]+x[2] == n, \
     ...:            combinations_with_replacement(range(1,n+1), 3)) 
     ...:            

In [385]: %timeit list(combinations_1(17))
5.9 ms ± 322 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [386]: %timeit list(combinations_2(17))
289 µs ± 28.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Upvotes: 0

JohnkaS
JohnkaS

Reputation: 631

In short, no, this is an extremely complex algorithm and would quickly sum up to O(n ^ 3)

The "algorithm" itself probably won't get a lot more efficient than O(n ^ 2), but you could change it to O(n ^ 2) fairly easily.

def combinations(n):
   for i in range(n - 2): # go up to n and add 1 + 1, assuming you don't want 0 and 0
       for j in range(n - 2): # the same again.
           if i + j >= n:
               continue
           yield (i, j, n - i - j) # there are probably more than just these, keep that in mind.

Hope this helps at least a bit.

Upvotes: 0

tobias_k
tobias_k

Reputation: 82899

You do not have to get all the combinations of three numbers. You can just get combinations of two numbers, and you know what the third number has to be.

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

This is much faster:

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

Upvotes: 4

Related Questions