Partition a number into a given set of numbers

Question

Here is what I am trying to do. Given a number and a set of numbers, I want to partition that number into the numbers given in the set (with repetitions). For example : take the number 9, and the set of numbers = {1, 4, 9}. It will yield the following partitions :

{ (1, 1, 1, 1, 1, 1, 1, 1, 1), (1, 1, 1, 1, 1, 4), (1, 4, 4), (9,)}

No other possible partitions using the set {1, 4, 9} cannot be formed to sum the number 9.

I wrote a function in Python which do the task :

S = [ 1, 4, 9, 16 ]
def partition_nr_into_given_set_of_nrs(nr , S):
   lst = set()
   # Build the base case :
   M = [1]*(nr%S[0]) + [S[0]] * (nr //S[0])
   if set(M).difference(S) == 0  :
      lst.add(M)
   else :
      for x in S :
         for j in range(1, len(M)+1):
            for k in range(1, nr//x +1 ) :
               if  k*x ==  sum(M[:j]) :
                  lst.add(  tuple(sorted([x]*k + M[j:])) )
   return lst

It works correctly but I want to see some opinions about it. I'm not satisfied about the fact that it uses 3 loops and I guess that it can be improved in a more elegant way. Maybe recursion is more suited in this case. Any suggestions or corrections would be appreciated. Thanks in advance.

Answer 1

I would solve this using a recursive function, starting with the largest number and recursively finding solutions for the remaining value, using smaller and smaller numbers.

def partition_nr_into_given_set_of_nrs(nr, S):
    nrs = sorted(S, reverse=True)
    def inner(n, i):
        if n == 0:
            yield []
        for k in range(i, len(nrs)):
            if nrs[k] <= n:
                for rest in inner(n - nrs[k], k):
                    yield [nrs[k]] + rest
    return list(inner(nr, 0))

S = [ 1, 4, 9, 16 ]
print(partition_nr_into_given_set_of_nrs(9, S))
# [[9], [4, 4, 1], [4, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1]]

Of course you could also do without the inner function by changing the parameters of the function and assuming that the list is already sorted in reverse order.

If you want to limit the number of parts for large numbers, you can add an aditional parameter indicating the remaining allowed number of elements and only yield result if that number is still greater than zero.

def partition_nr_into_given_set_of_nrs(nr, S, m=10):
    nrs = sorted(S, reverse=True)
    def inner(n, i, m):
        if m > 0:
            if n == 0:
                yield []
            for k in range(i, len(nrs)):
                if nrs[k] <= n:
                    for rest in inner(n - nrs[k], k, m - 1):
                        yield [nrs[k]] + rest
    return list(inner(nr, 0, m))