生成所有可能的n维k * k *…* k数组，每个数组沿轴均带有一线

Question

I'm working in Python and need to find an algorithm to generate all possible n-dimensional k,k,...,k-arrays each with a line of ones along an axis. So, the function takes two numbers - n and k and should return a list of arrays with all possible lines of k ones along each axis.

For example for n = 2 and k = 3 there are 6 possibilities (with 3 horizontal lines and 3 vertical):

[[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]]

For n = 3 and k = 3 there are 27 possibilities (9 lines with 3 ones each along each of 3 axes).

Unfortunately, I don't have even an idea of how to do it for arbitrary n and k. Any suggestions?

Answer 1

Here is a generator approach using itertools.product to get the indices for placing the line. itertools.product is often useful to replace nested loops of variable depth:

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

Without numpy, you can create the matrix recursively, and fill the ones along each of the n axis with a starting position that will vary in the k**(n-1) possibilities.

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

This might endup begin heavy (computations-wise), because there wil be n*k**(n-1) possibilities.