Search code examples
pythonalgorithmlist-comprehensionnested-lists

create list with elements which are sublists that alternate between elements in sublists from two other lists


The title is a bit weird but essentially I need to take two lists:

list1 = [['1', '1', '1', '1'], ['2', '2', '2', '2'], ['3', '3', '3', '3']]
list2 = [['a', 'a', 'a'], ['b', 'b', 'b']]

and then alternate between the elements in the sub-lists of the two lists and create sub-lists as elements in a new list whose sub-lists are the alternated elements from above

list3 = [['1', 'a', '1', 'a', '1', 'a', '1'], ['1', 'b', '1', 'b', '1', 'b', '1'], ['2', 'a', '2', 'a', '2', 'a', '2'], ... ]

right now my code is:

def foo(array1, array2):
    i = 1
    for sublist1 in array1:
        for sublist2 in array2:
            for val in sublist2:
                sublist1.insert(i, val)
                i += 2
        i = 1
    return array1

and I'm getting the output:

[['1', 'a', '1', 'a', '1', 'a', '1', 'b', 'b', 'b'], ['2', 'a', '2', 'a', '2', 'a', '2', 'b', 'b', 'b'], ['3', 'a', '3', 'a', '3', 'a', '3', 'b', 'b', 'b']]

the thing is that I'm working with much smaller lists as a proof of concept right now but the final algorithm will need to be able to do this for lists with millions of sub-lists.


Solution

  • I'd use itertools for the task:

    from itertools import product, zip_longest
    
    list1 = [["1", "1", "1", "1"], ["2", "2", "2", "2"], ["3", "3", "3", "3"]]
    list2 = [["a", "a", "a"], ["b", "b", "b"]]
    
    out = []
    for c in product(list1, list2):
        out.append([v for c in zip_longest(*c) for v in c if v is not None])
    
    print(out)
    

    Prints:

    [
        ["1", "a", "1", "a", "1", "a", "1"],
        ["1", "b", "1", "b", "1", "b", "1"],
        ["2", "a", "2", "a", "2", "a", "2"],
        ["2", "b", "2", "b", "2", "b", "2"],
        ["3", "a", "3", "a", "3", "a", "3"],
        ["3", "b", "3", "b", "3", "b", "3"],
    ]
    

    Note: If there's None in any of your sub-lists, use other fillvalue= in itertools.zip_longest