如何高效遍历元素范围受限、总和固定的指定长度非负整数数组

考虑所有长度为l、元素取值范围为0,...,m的非负整数数组，希望通过生成器仅遍历其中总和恰好等于s的数组。

例如，取l=7, s=5, m=4时，遍历结果样例如下：

(0, 0, 0, 0, 0, 1, 4)
(0, 0, 0, 0, 0, 2, 3)
(0, 0, 0, 0, 0, 3, 2)
(0, 0, 0, 0, 0, 4, 1)
(0, 0, 0, 0, 1, 0, 4)
(0, 0, 0, 0, 1, 1, 3)
(0, 0, 0, 0, 1, 2, 2)
(0, 0, 0, 0, 1, 3, 1)
(0, 0, 0, 0, 1, 4, 0)
[...]
(3, 2, 0, 0, 0, 0, 0)
(4, 0, 0, 0, 0, 0, 1)
(4, 0, 0, 0, 0, 1, 0)
(4, 0, 0, 0, 1, 0, 0)
(4, 0, 0, 1, 0, 0, 0)
(4, 0, 1, 0, 0, 0, 0)
(4, 1, 0, 0, 0, 0, 0)

不限制遍历的顺序，但要求算法足够高效。

以下是能实现上述需求但效率极低的暴力方案，当变量取值较大时运行速度过慢，无法满足要求：

import itertools
s = 5
l = 7
m = 5
for arr in itertools.product(range(m), repeat=l):
    if sum(arr) == s:
        print(arr)

采用带剪枝的回溯算法，每一步选择当前位置的取值时，直接根据剩余待分配总和、剩余位置数、元素上限三个条件限制可选值范围，完全不会生成无效组合，效率远高于暴力枚举。

代码实现 #

def generate_arrays(l: int, s: int, m: int):
    # 边界剪枝：参数本身无解时直接返回
    if s < 0 or l * m < s:
        return
    # 回溯递归逻辑
    def backtrack(pos, remaining, current):
        if pos == l:
            if remaining == 0:
                yield tuple(current)
            return
        # 当前位置可选最小值：保证剩余位置全取上限也能凑够剩余总和
        min_val = max(0, remaining - (l - pos - 1) * m)
        # 当前位置可选最大值：不超过元素上限，也不超过剩余需要凑的总和
        max_val = min(m, remaining)
        for val in range(min_val, max_val + 1):
            current.append(val)
            yield from backtrack(pos + 1, remaining - val, current)
            current.pop()
    yield from backtrack(0, s, [])

# 测试样例
if __name__ == "__main__":
    for arr in generate_arrays(7, 5, 4):
        print(arr)

效率说明 #

• 暴力方案时间复杂度为O(m^l)，只要l和m稍大就完全不可用
• 本方案时间复杂度仅和符合条件的数组总数正相关，无任何额外无效计算，是该问题下的最优解法
• 提前做了边界判断，输入参数本身无解时会直接返回，不会做无用计算

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