如何在Python中生成给定预算下的所有整数分箱组合？

我有一个预算N，希望将其拆分为K个分箱，限制条件为每个分箱只能包含整数，且所有分箱的数值之和等于预算。此类组合的数量为组合数C(N+K-1, K-1)。

例如当N=5、K=3时，我自己编写了如下代码：

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  

该代码返回结果：

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

请问是否存在可处理任意N和K的内置函数？我认为应该有，但目前尚未找到。

Python标准库中没有直接实现这个需求的内置函数，但可以借助itertools模块的工具高效实现，比嵌套循环的扩展性好得多。

核心思路是用「隔板法」：把N个预算看成N个相同的元素，要分成K个允许为空的分箱，相当于在N+K-1个间隔中选K-1个位置放置隔板，每个隔板之间的元素数量就是对应分箱的数值。用itertools.combinations可以直接生成所有隔板位置，再推导分箱数值：

import itertools

def generate_strategies(n, k):
    # 生成所有隔板位置的组合
    for partition_indices in itertools.combinations(range(n + k - 1), k - 1):
        prev_idx = -1
        strat = []
        for idx in partition_indices:
            # 计算当前分箱的数值：两个隔板之间的元素数
            strat.append(idx - prev_idx - 1)
            prev_idx = idx
        # 最后一个分箱的数值：从最后一个隔板到末尾的元素数
        strat.append((n + k - 1) - prev_idx - 1)
        yield strat

比如调用list(generate_strategies(5, 3))，得到的结果和你自己写的代码完全一致，而且支持任意正整数N和K。

如果不想用生成器，也可以直接返回列表：

def generate_strategies_list(n, k):
    strats = []
    for partition_indices in itertools.combinations(range(n + k - 1), k - 1):
        prev_idx = -1
        strat = []
        for idx in partition_indices:
            strat.append(idx - prev_idx - 1)
            prev_idx = idx
        strat.append((n + k - 1) - prev_idx - 1)
        strats.append(strat)
    return strats

这种方法的效率远高于嵌套循环，因为它直接生成符合条件的组合，不会遍历无效的数值组合。

内容的提问来源于stack exchange，提问作者Auriel