如何在Python中生成枚举元素的所有二维矩阵变体（生成器实现）

我定义了如下枚举类：

class MyEnum(enum.Enum):
   EMTPY = 0  # 注：原代码存在拼写错误，应为EMPTY
   APPLE = 1
   HONEY = 2

希望实现一个生成器函数variations(height, width)，传入高度和宽度参数，生成所有元素均为MyEnum枚举成员的二维数组变体。示例调用代码如下：

def variations(height, width):
    variation = # code
    yield variation

for v in variations(2,2):
    print(v)
    print("---------")

预期输出为所有可能的二维数组，例如：

[
  [EMTPY, EMTPY],
  [EMTPY, EMTPY]
]
---------
[
  [APPLE, EMTPY],
  [EMTPY, EMTPY]
]
---------
...

我尝试过itertools的组合、排列方法但未成功，因该任务比教程案例复杂。另外后续可能使用较大的height和width参数，所以必须以生成器形式实现。

可以利用itertools.product生成笛卡尔积，完美覆盖所有位置的枚举成员组合，配合生成器实现惰性迭代。

完整代码如下：

import enum
import itertools

class MyEnum(enum.Enum):
    EMPTY = 0  # 修正原拼写错误
    APPLE = 1
    HONEY = 2

def variations(height, width):
    enum_members = list(MyEnum)
    # 生成height*width个枚举成员的笛卡尔积（所有可能的一维组合）
    for flat_vals in itertools.product(enum_members, repeat=height * width):
        # 将一维结果切割为height行、width列的二维数组
        two_d_array = [
            flat_vals[i*width : (i+1)*width]
            for i in range(height)
        ]
        yield two_d_array

# 测试调用（格式化输出以提升可读性）
for arr in variations(2, 2):
    print("[")
    for row in arr:
        print(f"  [{', '.join(member.name for member in row)}],")
    print("]")
    print("---------")

代码说明 #

1. 笛卡尔积生成：itertools.product(enum_members, repeat=n)会生成所有长度为n的枚举成员组合，每个位置独立选择成员，正好对应二维数组展开后的所有可能情况。
2. 二维转换：将一维的笛卡尔积结果按width分割成height行，得到目标二维数组结构。
3. 生成器优势：itertools.product本身是惰性迭代的，配合yield实现生成器，不会一次性生成所有结果，即使参数较大也不会占用过多内存。

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