Search code examples
recursionmathcombinatoricsbacktrackingrecursive-backtracking

Algorithm to generate all nonnegative solutions to x1 + x2 + ....xk = n. (Stars and bars)

What is an algorithm to generate every solution to x1 + x2 + ....xk = n, xi >= 0? I am aware it is some form of backtracking but I am not sure how to implement this.

Solution

  • Without recursion, quicky modified from my answer with non-zero partitions, implements stars and bars principle

    from itertools import combinations

#print(parts(7,4))

def partswithzeros(n,r):
    result = []
    for comb in combinations(range(n+r-1), r-1):
        zpart = [comb[0]] + [comb[i]-comb[i-1]-1 for i in range(1,r-1)] + [n+r-2 - comb[-1]]
        result.append(zpart)
    return(result)

print(partswithzeros(5,3))

>>>[[0, 0, 5], [0, 1, 4], [0, 2, 3], [0, 3, 2], [0, 4, 1], [0, 5, 0], [1, 0, 4],
    [1, 1, 3], [1, 2, 2], [1, 3, 1], [1, 4, 0], [2, 0, 3], [2, 1, 2], [2, 2, 1],
    [2, 3, 0], [3, 0, 2], [3, 1, 1], [3, 2, 0], [4, 0, 1], [4, 1, 0], [5, 0, 0]]