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

Question

我在 Python 中工作，需要找到一种算法来生成所有可能的 n 维 k、k、...、k 数组，每个数组都有一行沿轴的数组。因此，该函数需要两个数字 - 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)

Answer 2

如果没有 numpy，您可以递归地创建矩阵，并在每个 n 轴上填充矩阵，其起始位置会因k**(n-1) 的可能性而异。

k=3
n=2
indexes = [0]*n
def build_zeroes(n, k):
    if n == 2:
       return [[0]*k for _ in range(k)]
    else:
       return [build_zeroes(n-1, k) for _ in range(k)]

def compute_coordinate(position, n, k):
    coords=[]
    for i in range(n):
        coords.append(position % k)
        position = position // k
    return coords

def set_in_matrix(m, coords, value=1):
    u = m
    for c in coords[:-1]:
        u = u[c]
    u[coords[-1]] = value

for axis in range(n):
    for start_position in range(k**(n-1)):
        coords = compute_coordinate(start_position, n-1, k)
        coords.insert(axis, 0)
        m = build_zeroes(n, k)
        for i in range(k):
            coords[axis] = i
            set_in_matrix(m, coords)
        print m

这最终可能会开始繁重（计算方面），因为会有n*k**(n-1) 的可能性。