Search code examples
pythongeneratorsequence-generators

Dynamically generating elements of list within list


I have a list, which is made up of the following elements,

list1 = [a1,a2,a3]

Where each element of this list can itself be a variable size list, eg,

a1 = [x1,y1,z1], a2 = [w2,x2,y2,z2], a3 = [p3,r3,t3,n3]

It's straight forward for me to set up a generator that loops through list1, and yields the constituents of each element;

array = []
for i in list1:
    for j in i:
        array.append[j]
        yield array

However, is there a way of doing this so I can specify the size of array?

eg - batch size of two;

1st yield : [x1,y1]
2nd yield : [z1,w1]
3rd yield : [x2,y2]
4th yield : [z2,p3]
5th yield : [r3,t3]
6th yield : [n3]
7th yield : repeat 1st

or batch size of 4;

1st yield : [x1,y1,z1,w1]
2nd yield : [x2,y2,z2,p3]
3rd yield : [r3,t3,n3]
4th yield : repeat first

It seems non-trivial to carry this out for different sized lists each containing other different sized lists inside.


Solution

  • This is pretty easy, actually, use itertools:

    >>> a1 = ['x1','y1','z1']; a2 = ['w2','x2','y2','z2']; a3 = ['p3','r3','t3','n3']
    >>> list1 = [a1,a2,a3]
    >>> from itertools import chain, islice
    >>> flatten = chain.from_iterable
    >>> def slicer(seq, n):
    ...     it = iter(seq)
    ...     return lambda: list(islice(it,n))
    ...
    >>> def my_gen(seq_seq, batchsize):
    ...     for batch in iter(slicer(flatten(seq_seq), batchsize), []):
    ...         yield batch
    ...
    >>> list(my_gen(list1, 2))
    [['x1', 'y1'], ['z1', 'w2'], ['x2', 'y2'], ['z2', 'p3'], ['r3', 't3'], ['n3']]
    >>> list(my_gen(list1, 4))
    [['x1', 'y1', 'z1', 'w2'], ['x2', 'y2', 'z2', 'p3'], ['r3', 't3', 'n3']]
    

    Note, we can use yield from in Python 3.3+:

    >>> def my_gen(seq_seq, batchsize):
    ...   yield from iter(slicer(flatten(seq_seq), batchsize), [])
    ...
    >>> list(my_gen(list1,2))
    [['x1', 'y1'], ['z1', 'w2'], ['x2', 'y2'], ['z2', 'p3'], ['r3', 't3'], ['n3']]
    >>> list(my_gen(list1,3))
    [['x1', 'y1', 'z1'], ['w2', 'x2', 'y2'], ['z2', 'p3', 'r3'], ['t3', 'n3']]
    >>> list(my_gen(list1,4))
    [['x1', 'y1', 'z1', 'w2'], ['x2', 'y2', 'z2', 'p3'], ['r3', 't3', 'n3']]
    >>>