从列表中生成循环批次

Question

基本上，我想从给定列表l中创建一个无限生成器，其批量大小为batch_size。举例来说，如果我有l = [1, 2, 3, 4, 5]和batch_size = 2名单，我想产生的[1, 2]，[3, 4]，[5, 1]，[2, 3]无限循环......（类似于itertool.circular使用附加批量大小）

我目前的做法是因为在结束后不给予正确的解决方案还没有我只是垫列表的第一个元素，当我到达终点：

l = [1, 2, 3, 4, 5] 

def generator(l, batch_size=2): 
    while 1: 
     for i in range(0, len(l), batch_size): 
      batch = l[i:(i + batch_size)] 
      if len(batch) < batch_size: 
       batch.extend(l[0: batch_size - len(batch)]) 
      yield batch 

>> gen = generator(l, batch_size=2) 
>> next(gen) # [1, 2] 
>> next(gen) # [3, 4] 
>> next(gen) # [5, 1] 
>> next(gen) # [1, 2] << want [2, 3] instead 

有没有办法做到这一点以循环方式？

Answer 1

是的，你基本上要的 “取” 和cycle组合：

>>> def circle_batch(iterable, batchsize): 
...  it = itertools.cycle(iterable) 
...  while True: 
...   yield list(itertools.islice(it, batchsize)) 
... 
>>> l = [1, 2, 3, 4, 5] 
>>> c = circle_batch(l, 2) 
>>> next(c) 
[1, 2] 
>>> next(c) 
[3, 4] 
>>> next(c) 
[5, 1] 
>>> next(c) 
[2, 3] 
>>> next(c) 
[4, 5] 

从recipes in the docs你会看到take是一个基本的工具，因此使用：

>>> def take(n, iterable): 
...  "Return first n items of the iterable as a list" 
...  return list(islice(iterable, n)) 
... 
>>> def cycle_batch(iterable, batchsize): 
...  it = itertools.cycle(iterable) 
...  while True: 
...   return take(batchsize, it) 
... 
>>> l = [1, 2, 3, 4, 5] 
>>> c = circle_batch(l, 2) 
>>> next(c) 
[1, 2] 
>>> next(c) 
[3, 4] 
>>> next(c) 
[5, 1] 
>>> next(c) 
[2, 3] 
>>> next(c) 
[4, 5] 
>>> next(c) 
[1, 2] 

Answer 2

这应该工作：

def generator(l, batch_size=2): 
    gen = iter(itertools.cycle(l)) 
    while 1: 
     yield [next(gen) for _ in range(batch_size)] 

gen = generator(l, batch_size=2)