How to create all possible integer binning combinations given a budget in Python?-CodePudding

I have a budget N, which I would like to split into K bins, with the restriction being that each bin can only contain integers and the sum of each bin must equal to the budget. As such, I want to have (N K-1)C(K-1) possible combinations, using nCr notation.

For example, for N = 5, K = 3, I have the following code:

def generateStrategies():
    s = 5
    strats = []
    for x1 in range(s   1):
        current_strat = []
        for x2 in range((s   1) - x1):
            for x3 in range((s   1) - (x1   x2)):
                if (x1   x2   x3) == s:
                    current_strat = [x1,x2,x3]
                    strats.append(current_strat)
    return strats

which returns

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

Is there an inbuilt function that does this for any N and K? I feel like there probably is, but so far I can't find it.

CodePudding user response：

It isn't builtin, but I think this is what you are asking for:

def budgetSplit(n,k,sofar=[]):
    if k==1:
        yield sofar   [n]
    else:
        for n2 in range(0,n 1):
            for b in budgetSplit(n-n2,k-1,sofar [n2]):
                yield b 

print(list(budgetSplit(5,4)))

CodePudding user response：

from itertools import combinations as c
list(c([0,1,2,3,4,5],3))

list(c(list(range(5 1)),3))