python - 生成所有可能的 n 维 k*k*...*k 数组，每个数组都有沿轴的行

Question

我在 Python 中工作，需要找到一种算法来生成所有可能的 n 维 k,k,...,k 数组，每个数组都沿轴有一行 1。因此，该函数接受两个数字 - n 和 k，并且应该返回一个数组列表，其中包含沿每个轴的所有可能的 k 行。

例如，对于 n = 2 和 k = 3，有 6 种可能性(3 条水平线和 3 条垂直线):

[[1, 1, 1], 
 [0, 0, 0], 
 [0, 0, 0]],
[[0, 0, 0], 
 [1, 1, 1], 
 [0, 0, 0]],
[[0, 0, 0], 
 [0, 0, 0], 
 [1, 1, 1]],
[[1, 0, 0], 
 [1, 0, 0], 
 [1, 0, 0]],
[[0, 1, 0], 
 [0, 1, 0], 
 [0, 1, 0]],
[[0, 0, 1], 
 [0, 0, 1], 
 [0, 0, 1]]

对于 n = 3 和 k = 3，有 27 种可能性(9 条线，沿 3 个轴各有 3 个)。

不幸的是，我什至不知道如何为任意 n 和k 做这件事。有什么建议吗？

Answer 1

这是一个生成器方法，使用 itertools.product 来获取用于放置线条的索引。 itertools.product 通常用于替换可变深度的嵌套循环:

import numpy as np
import itertools

def lines(n, k):
    for axis in range(n):
        ranges = ((slice(None),) if a==axis else range(k) for a in range(n))
        for idx in itertools.product(*ranges):
            ret = np.zeros(n*(k,), dtype=int)
            ret[idx] = 1
            yield ret

for line in lines(2, 3):
    print(line)