Python实现两个NumPy数组的所有唯一组合生成、存储及CSV导出

1. 核心逻辑说明 #

你需要的343种组合本质是y的3次可放回抽样有序组合，每个组合对应2行3列的结构：第一行固定为x的变量名，第二行是对应取值，所有组合可以存入形状为(343, 2, 3)的numpy数组中。

2. 完整代码实现 #

首先导入依赖库：

import numpy as np
import itertools
import csv

生成所有组合并存入大数组 #

# 定义原始数组
x = np.array(['var1','var2','var3'])
y = np.array([1,2,3,4,5,6,7])

# 生成所有取值组合（可重复有序，共7^3=343组）
value_combinations = list(itertools.product(y, repeat=len(x)))

# 初始化存储所有2行3列数组的大数组，dtype用object兼容字符串和数字类型
all_result_arrays = np.empty((len(value_combinations), 2, 3), dtype=object)

for idx, vals in enumerate(value_combinations):
    # 垂直堆叠变量名行和取值行，得到单组结果
    all_result_arrays[idx] = np.vstack([x, np.array(vals)])

你可以直接通过索引访问任意一组结果，比如all_result_arrays[0]就会返回第一组的2行3列数组：

[['var1' 'var2' 'var3']
 [1 1 1]]

纯numpy生成组合的替代方案 #

如果不想引入itertools依赖，可以用numpy的meshgrid实现同样的效果：

y1, y2, y3 = np.meshgrid(y, y, y, indexing='ij')
value_combinations = np.stack([y1.ravel(), y2.ravel(), y3.ravel()], axis=1)

3. 写入CSV文件 #

提供两种写入格式，根据你的需求选择：

格式1：常规结构化CSV（适合后续数据分析） #

第一行是变量名表头，后续每一行对应一组取值，文件共344行3列：

with open('combinations_structured.csv', 'w', newline='', encoding='utf-8') as f:
    writer = csv.writer(f)
    # 写入表头
    writer.writerow(x)
    # 写入所有取值组合
    writer.writerows(value_combinations)

格式2：保留2行3列的原始结构 #

每组结果单独占2行，组间用空行分隔，完全匹配你要求的展示形式：

with open('combinations_raw_format.csv', 'w', newline='', encoding='utf-8') as f:
    writer = csv.writer(f)
    for arr in all_result_arrays:
        writer.writerows(arr)
        # 组间插入空行分隔
        writer.writerow([])

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