CODEWITHSUNDEEP
pythonarraysnumpy
Chad Larson
Chad Larson

Reputation: 309

Iterate over all possible numpy binary arrays with restrictions using python

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?

Upvotes: 0

Views: 1226

Answers (2)

HYRY
HYRY

Reputation: 97331

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

Upvotes: 1

pyan
pyan

Reputation: 3707

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]

Upvotes: 1

Related Questions