遍歷所有可能的numpy二進制數組限制使用python

Question

我想遍歷所有可能的大小爲mxn的二進制數組，但有一些限制。正如你所知，隨着m和n的增加（2 ^（m * n）數組），該組數組變得極端。我已經寫了一些代碼來遍歷所有這些可能性。

mxn = np.arange(m*n).reshape(m,n) 
for i in xrange(0, 2**(m*n)): 
    arr = (i >> mxn) % 2 
    print arr 

我可以添加一些額外的限制，減少我需要迭代的數組。限制是矩陣中的每一行可以不超過1.第二個限制是矩陣中所有元素的總和不能大於m。我可以修改我已經完成的任務嗎？還是有一條完全不同的道路，我應該走下去？

Answer 1

由於矩陣的大小爲m*n，如果滿足第一個限制，則會自動滿足第二個限制。

由於每行至多有1個元素不爲零，所以一行只有n+1選項。假設有m行，這種矩陣的組合的可能數量因此是(n+1)**m。下面的代碼將複雜度從2**(m*n)降低到(n+1)**m。

rows = [] 
r = np.zeros((n,), dtype=np.int) 
rows.append(r) 
for i in range(n): 
    r = np.zeros((n,), dtype=np.int) 
    r[i] = 1 
    rows.append(r) 

rows = np.array(rows) 

arr = np.zeros((m,n), dtype=np.int) 
idx = [0]*m 
base = np.array([(n+1)**j for j in range(m)]) 
for i in xrange(0, (n+1)**m): 
    idx = (i/base) % (n+1) 
    print rows[idx] 

Answer 2

相同的算法@pyan，但使用numpy的功能：

import numpy as np 

m, n = 2, 3 
rows = np.vstack([np.zeros(n, dtype=np.int), np.identity(n, dtype=np.int)]) 
index = np.indices([len(rows)] * m).reshape(m, -1).T 
rows[index] 

輸出是一個三維數組：

[[[0 0 0] 
    [0 0 0]] 

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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