能否用Python标准库itertools或numpy生成指定结构的序列数据？

问题 #

我编写了如下函数e(w)：

def e(w):
    for i in range(w):
        for j in range(w-1):
            for k in range(w-2):
                yield [i,j,k]

调用print([i for i in e(4)])可生成如下格式的数据：

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

随后我将其泛化为递归函数f()：

def f(w, mx=0, depth=0, values=[]):
    if mx==0:
        yield values
    for i in range(w - depth):
        yield from f(w, mx-1, depth+1, values + [i])

调用print([i for i in f(4,3)])可生成相同数据。请问能否使用itertools或numpy等Python标准库实现该数据的生成？

使用Python标准库实现的方案 #

用itertools实现 #

可以利用itertools.product生成多维度笛卡尔积，匹配原函数的取值规则：

import itertools

def generate_with_itertools(w, length=3):
    # 为每个位置生成对应的取值范围
    ranges = [range(w - d) for d in range(length)]
    # 生成所有组合并转为列表格式
    return [list(item) for item in itertools.product(*ranges)]

# 调用示例
print(generate_with_itertools(4, 3))

逻辑和原函数完全一致：第1个元素取自range(w)，第2个取自range(w-1)，第3个取自range(w-2)，通过product直接生成所有可能的组合。

用numpy实现 #

如果需要numpy数组格式的结果，可通过meshgrid构建网格后变形得到：

import numpy as np

def generate_with_numpy(w, length=3):
    # 生成每个维度的取值数组
    axes = [np.arange(w - d) for d in range(length)]
    # 创建网格（indexing='ij'保证顺序和笛卡尔积一致）
    grids = np.meshgrid(*axes, indexing='ij')
    # 展开每个网格并拼接成二维数组
    result = np.column_stack([grid.ravel() for grid in grids])
    # 如需列表格式则转成list，否则直接返回numpy数组
    return result.tolist()

# 调用示例
print(generate_with_numpy(4, 3))

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