复制生成器而不会耗尽内存

时间:2016-01-22 23:02:35

标签: python generator itertools

我正在编写一个python类,它会找到所有可能的magic squares给定一个整数size和一个可能combinations的生成器。这些组合是长度为size**2的元组,并分为size×size网格。代码本身工作正常,但重用生成器似乎需要itertools.tee。在下面显示的示例中,这会导致线程使用的内存跳转到300MB,因为迭代器中的每个值都存储在列表中。

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

有一个明显的解决方案可以解决这个问题。首先,我可以要求提供生成器而不是请求生成器本身的函数,但这需要用户包装它们的生成器表达式。其次,我可以缓存解决方案。为了论证,假设我不再需要检查对角线,如果没有足够数量的解决方案,那么我需要更新checkGrid并重申combinations

所以,我的问题是:如果没有创建这个潜在的巨大内存问题,是否真的没有办法复制生成器?我不关心保留生成器的部分状态,我只是想让它迭代与原始生成器相同的值。

修改

在Python 3.X中,您可以使用copy.deepcopy来复制依赖项都可以选择的itertools个对象。

4 个答案:

答案 0 :(得分:1)

没有什么是不可能的......

以下情况适用于itertools.permutations。不要认为它适用于任何迭代,因为它不会!

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

答案 1 :(得分:0)

无法复制任意迭代器。极少数特定的迭代器类型支持复制;我唯一知道的是itertools.tee。但是,一般情况下,迭代器可能有太多不可复制的依赖关系,因此复制机制是迭代器协议的一部分。

您只是遇到了这个问题,因为您已经编写了一个试图获取一次性迭代器并返回一个可重用对象的API。如果您要使用迭代器,则应设计API以返回迭代器而不是可以创建一次的MagicSquare对象,然后重复调用solutions

对于您的用例,我建议将MagicSquare作为生成器。主要的,可能只用于此类似乎是为解决方案的迭代器调用solutions。为什么不简单地用一个执行MagicSquare(size, combinations).solutions()当前功能的函数替换类?

答案 2 :(得分:0)

不是传递生成器,而是传递一个函数,该函数在调用时返回一个新的生成器。这将允许MagicSquare根据需要多次迭代组合,而不会将它们保留在内存中。

用你的代码来解释:

class MagicSquare:
    def __init__(self, size, get_combinations):
        ...
        self.get_combinations = get_combinations

    ...

    def solutions(self):
        for entries in self.get_combinations():
            ...

if __name__ == '__main__':
    combs2 = lambda: permutations(range(20,30), 9) # 
    ms2 = MagicSquare(3, combs2)
    ...

答案 3 :(得分:-1)

由于您的生成器是独立且确定的,因此使用两个副本的最佳方法是创建其中两个副本。 (如有必要,修改MagicSquare的签名以接受两个生成器;但看起来您希望将该副本用于其他目的?)

combs2a = permutations(range(20,30), 9)
combs2b = permutations(range(20,30), 9)