ph0t3k ph0t3k - 1 month ago 11
Python Question

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?

Answer

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.