CODEWITHSUNDEEP
python
user1444165
user1444165

Reputation:

Generating products of lists that sum to a parameter N

I'm looking for a fast way to generate products of values 0 to (and including) N, that sum to N. For simpler problems, this looks like

N = 10
output = []
for i in xrange(N+1):
    for j in xrange(N+1-i):
        for k in xrange(N+1-i-j):
           output.append([i,j,k])

So, output should only contain lists of length 3, in this case. Similar to

[v for v in itertools.product(*[range(N+1) for i in range(3)]) if sum(v) == N)]

EDIT: As Joran noted, the lines above do not take advantage of the fact that the third element will always be the 'rest', i.e. N-i-j.

Hopefully, a way to do this with itertools without evaluating the outcome of every possible product exists. Currently, I have a recursive function that, apart from some efficiency issues (often calling len for example), takes quite some time for large max_depth:

def recursive_expand(input_list, max_sum, max_depth):
    ''' Creates list of lists of integers, where elements of each 
        sub_list sum to max_sum.
    '''
    # If current elements already sum to max_sum
    if max_sum == 0:
        return [input_list + [0] * (max_depth - len(input_list))]
    # If the list can only contain one more element
    elif len(input_list) == max_depth - 1:
        return [input_list + [max_sum]]

    output_lists = []
    for n in xrange(max_sum + 1):
        output_lists.extend( \
                 recursive_expand(input_list + [n], max_sum - n,max_depth))
    return output_lists

>>> foo = recursive_expand([], 2, 3)
>>> print np.array(foo)
[[0 0 2]
 [0 1 1]
 [0 2 0]
 [1 0 1]
 [1 1 0]
 [2 0 0]]

EDIT 2: Credit to jonrsharpe, this function generates tuples of length k containing integers from start to n, summing to n, where ordering matters. See his answer if ordering does not matter (as it is more efficient).

def sets_2(n, k, start):
    for x in xrange(start, n + 1):
        if k == 2:
            yield x, n-x
        elif n-x == 0:
            yield (x,) + (0,) * (k-1)
        else:
            for tup in sets(n-x, k-1, x):
                yield (x,) + tup

And, turns out the most efficient implementation for finding integer partitions (which can be padded with zeros afterwards) can be found here: http://homepages.ed.ac.uk/jkellehe/partitions.php

Upvotes: 2

Views: 98

Answers (2)

jonrsharpe
jonrsharpe

Reputation: 122032

If you aren't interested in the order (i.e. you always want x <= y <= z, only unique sets of integers), this should be reasonably efficient, by:

  1. limiting the range you search for appropriate x and y;
  2. calculating z from x and y, rather than looping again; and
  3. yielding the values to create a generator.

Code:

def sets(n):
    for x in xrange(1, (n // 3) + 1):
       for y in xrange(x, ((n - x) // 2) + 1):
           yield x, y, n - (x + y)

For example:

list(sets(10)) == [(1, 1, 8), (1, 2, 7), (1, 3, 6), (1, 4, 5), 
                   (2, 2, 6), (2, 3, 5), (2, 4, 4), (3, 3, 4)]

For arbitrary set length k, you could use recursion:

def sets_2(n, k=3, s=1):
    for x in range(s, (n // k) + 1):
        if k == 2:
            yield x, n - x
        else:
            for s in sets_2(n - x, k - 1, x):
                yield (x, ) + s

For example:

list(sets_2(10, 2)) == [(1, 9), (2, 8), (3, 7), (4, 6), (5, 5)]

list(sets_2(10, 4)) == [(1, 1, 1, 7), (1, 1, 2, 6), (1, 1, 3, 5), 
                        (1, 1, 4, 4), (1, 2, 2, 5), (1, 2, 3, 4), 
                        (1, 3, 3, 3), (2, 2, 2, 4), (2, 2, 3, 3)]

To include zeros, set s=0:

list(sets_2(10, 3, 0)) == [(0, 0, 10), (0, 1, 9), (0, 2, 8), (0, 3, 7), 
                           (0, 4, 6), (0, 5, 5), (1, 1, 8), (1, 2, 7), 
                           (1, 3, 6), (1, 4, 5), (2, 2, 6), (2, 3, 5), 
                           (2, 4, 4), (3, 3, 4)]

Upvotes: 2

Joran Beasley
Joran Beasley

Reputation: 113988

OK Im going to go out on a limb here and guess that you want 3 values that sum to some magic value N, while also meeting some additional constraint

you dont need to compute all 3 values

a + b +c = N  ==> c = N - a -b

you can compute 2 values and solve the 3rd

for a in range(N): #unless the other 2 values can be 0 you dont really need to go all the way to n
    for b in range(a,N):  #this assumes that b >= a
         c = N - a - b
         if some_constraint(a,b,c,N):
                my_list.append([a,b,c])

to be honest I have no idea if this answers your question, however it seems like this is the direction your question was heading ...

Upvotes: 1

Related Questions