I'm writing a python class that finds all possible magic squares given an integer size and a generator for the possible combinations. These combinations are tuples of length size**2 and are split into a size×size grid. The code itself works fine, but reusing the generator appears to require itertools.tee. In the example shown below, this causes the memory used by the thread to jump to 300MB since every value in the iterator is being stored in a list.
from itertools import permutations, tee
class MagicSquare:
def __init__(self, size, combinations):
self.size = size
self.range = range(self.size)
self.combinations = combinations
def getGrid(self, entries):
return [ entries[self.size*i:self.size*(i+1)] for i in self.range ]
def checkGrid(self, grid):
check_sum = sum(grid[0])
if any( sum(row) != check_sum for row in grid ):
return False
if any( sum(row[col] for row in grid) != check_sum for col in self.range ):
return False
if sum(grid[diag][diag] for diag in self.range) != check_sum:
return False
if sum(grid[diag][self.size-diag-1] for diag in self.range) != check_sum:
return False
return True
def solutions(self):
combinations, self.combinations = tee(self.combinations)
for entries in combinations:
grid = self.getGrid(entries)
if self.checkGrid(grid):
yield grid
if __name__ == '__main__':
combs = permutations(range(20,30), 9)
ms = MagicSquare(3, combs)
for solution in ms.solutions():
for row in solution:
print row
print
There are two obvious solutions to this issue that come to mind. First, I could ask for a function that provides a generator rather than ask for a generator itself, but this requires the user to wrap their generator expression. Second, I could cache the solutions. For the sake of argument, let's say I no longer want to check the diagonals if there aren't a sufficient number of solutions, so I'll need to update checkGrid and reiterate over combinations.
So, my question is: Is there really no way to copy a generator without creating this potentially huge memory issue? I don't care about retaining a partial state of the generator, I just want it to iterate over the same values as the original generator.
It looks like in Python 3.X, you can use copy.deepcopy to copy itertools objects whose dependencies are all pickable.
Nothing is impossible...
The following happens to work well with itertools.permutations. Don't assume that it will work with any iterable, because it won't!
>>> from itertools import permutations
>>> combs = permutations(range(20,30), 9)
>>> from copy import deepcopy
>>> combs2 = deepcopy(combs)
>>> next(combs)
(20, 21, 22, 23, 24, 25, 26, 27, 28)
>>> next(combs2)
(20, 21, 22, 23, 24, 25, 26, 27, 28)