如何高效生成满足和≤t的d维非负整数数组？Python实现

一、借助Itertools实现的方案 #

1. 基础（但低效）的暴力法 #

对于较小的d和t，可以直接用itertools.product生成所有可能的数组，再筛选出元素和≤t的结果。示例代码：

import itertools

d = 2
t = 2

all_arrays = itertools.product(range(t+1), repeat=d)
valid_arrays = [list(arr) for arr in all_arrays if sum(arr) <= t]

print(valid_arrays)
# 输出: [[0,0], [0,1], [0,2], [1,0], [1,1], [2,0]]

但这种方法时间复杂度为O(t^d)，当d较大（比如d≥10、t≥5）时，会产生大量无效组合，效率极低，不推荐使用。

2. 基于组合数学的高效实现（推荐） #

利用星与条定理的变形：将问题转化为求解方程x₁ + x₂ + ... + x_d + x_{d+1} = t的非负整数解（其中x_{d+1}是松弛变量，代表未使用的额度）。每个解对应原问题中一个元素和≤t的数组（只需忽略x_{d+1}）。

生成所有解的核心思路：

• 把t个“星”和d个“条”（分隔d+1个变量）排列，共有C(t + d, d)种排列方式
• 用itertools.combinations生成条的位置，再计算每个变量的取值

示例代码：

import itertools

def generate_valid_arrays(d, t):
    # 生成d个分隔符在t+d个位置中的索引（范围0到t+d-1）
    for dividers in itertools.combinations(range(t + d), d):
        # 计算每个元素的值：相邻分隔符的距离减1
        arr = [dividers[0]]
        for i in range(1, d):
            arr.append(dividers[i] - dividers[i-1] - 1)
        yield arr

# 测试示例
d = 2
t = 2
valid_arrays = list(generate_valid_arrays(d, t))
print(valid_arrays)
# 输出: [[0,0], [0,1], [0,2], [1,0], [1,1], [2,0]]

这种方法时间复杂度为O(C(t+d, d))，远低于暴力法，尤其适合d较大的场景。比如d=10、t=5时，暴力法要生成6^10≈600万组，而高效法仅生成C(15,10)=3003组，效率提升显著。

二、进一步优化思路 #

如果d和t都很大，生成所有数组可能占用大量内存，建议使用生成器（如代码中的yield）逐个返回结果，避免一次性加载所有数据到内存。

另外，若仅需统计符合条件的数组数量而非生成具体数组，直接用组合数公式C(t + d, d)计算即可，无需生成任何数组。

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