使用python遍历所有可能的numpy二进制数组

Question

I'm trying to iterate over all possible binary arrays of size mxn but with some restrictions. As you know the set of arrays becomes extreme as m and n increase (2^(m*n) arrays). I have written some code that will iterate over all of these possibilities.

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

I can add some additional restrictions that reduce the arrays I need to iterate over. The restrictions are that each row in the matrix can sum to no more than 1. The second restriction is that the sum of all of the elements in the matrix cannot be greater than m. Can I modify what I already have to accomplish this, or is there an altogether different path that I should be going down?

Answer 1

Since the matrix is of size m*n , the second restriction will be automatically met if the first one is satisfied.

Since each row has at most 1 element to be none zero, there are only n+1 choices for a row. Given there is m rows, the possible number of combinations for such a matrix is thus (n+1)**m . The following code then reduces the complexity from 2**(m*n) to (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

The same algorithm as @pyan, but use numpy functions:

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]

output is a 3d array:

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