Iterator用于可变数量的嵌套循环

时间:2013-12-07 12:12:13

标签: python numpy iterator

我正在尝试为这样的事情编写一个迭代器(从原始问题中简化):

import numpy as np

class IterableNumber:
  def __init__(self, digits):
    self.digits = digits
    self.zero_indices = np.where(digits == 0)[0]
  def __iter__(self):
    return self
  def next(self):
    # For each index in self.zero_indices, loop 1 through 9?
    pass
  # Other methods omitted

IterableNumber采用数组,但“0”是通配符。当我遍历这个对象时,我希望每个'0'数字循环1到9的范围,即

num = IterableNumber(np.array([5, 5, 0, 5, 0]))
for digits in num:
  print digits

应该打印

[5, 5, 1, 5, 1]
[5, 5, 1, 5, 2]
...
[5, 5, 1, 5, 9]
[5, 5, 2, 5, 1]
[5, 5, 2, 5, 2]
...
...
[5, 5, 9, 5, 9]

为一般情况编写这个迭代器似乎需要递归才能处理可变数量的嵌套循环(每个存在一个0),就像这个站点上的其他线程所暗示的那样。但是如何在迭代器的next()函数的上下文中编写这样的递归函数呢?或者可能有另一种方法可以完全解决这个问题?任何见解将不胜感激:)

如果它是相关的,我在python 2.7.3上。

1 个答案:

答案 0 :(得分:3)

您可以使用itertools.product

import numpy as np

from itertools import product

class IterableNumber:
    def __init__(self, digits):
        self.digits = digits
        self.zero_indices = np.where(self.digits==0)[0]
        self.length = len(self.zero_indices)

    def __iter__(self):
        for x in product(range(10), repeat=self.length):
            self.digits[self.zero_indices] = np.array(x)
            yield self.digits

<强>演示:

>>> for x in IterableNumber(np.array([5, 5, 0, 5, 0])):
    print x
...     
[5 5 0 5 0]
[5 5 0 5 1]
[5 5 0 5 2]
[5 5 0 5 3]
[5 5 0 5 4]
[5 5 0 5 5]
[5 5 0 5 6]
[5 5 0 5 7]
[5 5 0 5 8]
[5 5 0 5 9]
...
...
[5 5 9 5 3]
[5 5 9 5 4]
[5 5 9 5 5]
[5 5 9 5 6]
[5 5 9 5 7]
[5 5 9 5 8]
[5 5 9 5 9]
>>>