CODEWITHSUNDEEP
pythonpython-itertools
ph0t3k
ph0t3k

Reputation: 5

avoid computing whole cartesian product (python itertools)

My goal is to find any permutation of a diagonal in a nxn matrix (2 <= n <= 15). The matrix consists of zeros and ones.

Currently I do it like this:

  indices = [[j for j, x in enumerate(row) if x == 1] 
                for row in self.matrix]
  cart = list(itertools.product(*indices))
  cart = [list(tup) for tup in cart]
  cart = filter(lambda dia: len(list(set(dia))) == len(dia), cart)
  return cart

This works fine if the matrix is not too large, but otherwise it fails with: MemoryError

So is there a way to avoid the whole computation of cart? so that it for example one permutation is found, the computation stops?

Upvotes: 0

Views: 605

Answers (2)

John Zwinck
John Zwinck

Reputation: 249153

You can simplify the last part of your code to have it just return the first answer:

def foo(matrix):
    indices = [[j for j, x in enumerate(row) if x == 1] for row in matrix]

    # this part is changed, very simple and efficient now
    for dia in itertools.product(*indices):
        if len(set(dia)) == len(dia):
            return dia

In other words, don't be so clever with filter and lambda and all that--it's unnecessary.

Upvotes: 0

Moses Koledoye
Moses Koledoye

Reputation: 78556

Simply make all the evaluations lazy by not calling list on the result of itertools.product and using itertools.ifilter in place of filter:

from itertools import ifilter, product

indices = [[j for j, x in enumerate(row) if x == 1]  for row in self.matrix]
cart = product(*indices)
found_cart = next(ifilter(lambda dia: len(set(dia)) == len(dia), cart), None)

next returns the first case where the predicate in ifilter is True or returns None in case there is no matching item.

The computation stops once a matching item is found.

Upvotes: 1

Related Questions