求解和为n的定长M整数组合的递归/itertools实现方案

给定非负整数n与正整数M，输出所有长度为M的非负整数序列，满足序列内所有元素之和等于n。

• 示例输入：n=2、M=3
• 示例输出：[2,0,0]、[0,2,0]、[0,0,2]、[1,1,0]、[1,0,1]、[0,1,1]
• 要求：方案兼容任意合法取值的n和M，优先提供基于itertools的实现。

该问题对应组合数学中「求方程x₁+x₂+…+x_M = n的所有非负整数解」的经典场景，等价于n个相同元素放入M个有序盒子的全部分法。基于隔板法原理，结合itertools.combinations枚举隔板位置即可直接生成所有合法序列，无无效遍历，效率远高于暴力枚举所有可能值再筛选和的方案。

import itertools

def sum_sequences(n: int, M: int) -> list[list[int]]:
    result = []
    # 隔板法总长度为n+M-1，选M-1个隔板位置
    for splits in itertools.combinations(range(n + M - 1), M - 1):
        current_seq = []
        prev_pos = -1
        for split_pos in splits:
            # 两个隔板之间的元素个数即为对应位置的数值
            current_seq.append(split_pos - prev_pos - 1)
            prev_pos = split_pos
        # 补全最后一段的数值
        current_seq.append(n + M - 1 - prev_pos - 1)
        result.append(current_seq)
    return result

# 示例验证
if __name__ == "__main__":
    print(sum_sequences(n=2, M=3))

运行上述示例代码，将输出符合要求的全部序列：

[[2, 0, 0], [1, 1, 0], [1, 0, 1], [0, 2, 0], [0, 1, 1], [0, 0, 2]]

注：结果顺序由组合生成逻辑决定，若需要特定排序规则，可在返回结果前按需排序。

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