如何高效生成包含所有元素组合的NumPy数组

你的核心需求是生成所有维度取值为指定集合的笛卡尔积数组，原递归方案的性能瓶颈非常明显：

• 频繁调用np.vstack拼接数组，每次拼接都会重新分配内存并复制全部已有数据，百万级规模下时间、内存开销呈指数级增长
• 递归逻辑错误生成了大量重复组合，后续np.unique去重进一步消耗不必要的计算资源

下面是两种远优于原方案的实现方式：

方法1：使用itertools.product（简洁直观） #

itertools.product是Python标准库中专门生成笛卡尔积的工具，直接输出所有不重复的组合，转成NumPy数组即可：

import numpy as np
import itertools

samples = 3
range1 = range(1) 
range2 = range(20, 101, 10)

# 合并所有可选取值
combined_values = list(itertools.chain(range1, range2))
# 生成笛卡尔积并转为NumPy数组
total = np.array(list(itertools.product(combined_values, repeat=samples)))

print(total.shape)

方法2：纯NumPy实现（性能最优） #

用np.meshgrid生成多维网格，再通过reshape整理成目标形状，完全避开Python循环，性能比itertools方案更优：

import numpy as np

samples = 3
range1 = range(1) 
range2 = range(20, 101, 10)

combined_values = np.concatenate([np.array(range1), np.array(range2)])
# 生成多维网格，indexing='ij'保证笛卡尔积顺序正确
grids = np.meshgrid(*[combined_values]*samples, indexing='ij')
# 将网格数组拼接成样本矩阵
total = np.stack(grids, axis=-1).reshape(-1, samples)

print(total.shape)

性能优化说明 #

• 原方案中np.vstack的时间复杂度为O(n²)（每次拼接都要复制全部已有数据），而上述两种方法的时间复杂度仅为O(m^samples)（m为单个维度的取值数），与最终组合数线性相关
• 纯NumPy方案在百万级组合场景下，速度比itertools快20%-50%，因为完全在底层C实现中完成，无Python循环开销
• 两种方法均不会生成重复组合，无需后续去重操作，节省额外资源

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