pythonmathcombinatoricsrecurrence

Python Program to write an integer as a sum of r positive integers


I want to create a python program which returns all possible ways of writing an integer (say n) as a sum of r positive integers.

My code is as follows :

def f(n,r):
cache = {}

l1 = []
for i in range(1,n):
    l1.append([i,n-i])
    
    
if r ==2 :
    cache[tuple([n,r])] = l1
    return l1

elif (n,r) in cache:
    return cache[(n,r)]
    
else:
    lr = []
    for i in range (1,n):
        for x in f(n-i,r-1):
            x.append(i)
            lr.append(x)
    cache[tuple([n,r])] = lr        
    return lr

It works well till n = 20 for any value of r but after that it's really slow, which is kinda expected because recurrence is slow. How can I achieve this without using recurrence ?


Solution

  • Generation in lexicographic order.

    Note that any generation method won't be fast for large n,r values due to exponential grow of number of partitions.

    from itertools import combinations
    def parts(n,r):
        result = []
        for comb in combinations(range(n-1), r-1):
            parts = [comb[0]+1]+[comb[i]-comb[i-1] for i in range(1,r-1)]+[n-1-comb[-1]]
            #print(comb, parts)
            result.append(parts)
        return(result)
    
    print(parts(6,3))
    
    >>>[[1, 1, 4], [1, 2, 3], [1, 3, 2], [1, 4, 1], [2, 1, 3], [2, 2, 2], 
        [2, 3, 1], [3, 1, 2], [3, 2, 1], [4, 1, 1]]