Search code examples
pythonfor-loopnested-listsiterable

Iterating through non-rectangular nested lists of objects column by column using for-loop


I have a situation where I'm given a nested list that has very good odds of being non-rectangular. For example:

lists = [[obj1,obj2,obj3],[obj4,obj5],[obj6,obj7,obj8,obj9]]

and I want to iterate over it such that x = obj1,obj4,obj6,obj2,obj5,obj7,obj3,obj8,0bj9. Iterating over the elements of the first column, then the elements of the second column, and so on.

It could easily be done by transposing it and iterating over it, but I've yet to figure out how to do that.


Solution

  • Itertools is really handy for this. You can combine zip_longest chain.from_iterable and filter to make an iterator that gives you the result you're looking for:

    from itertools import zip_longest, chain
    
    lists = [['obj1','obj2','obj3'],['obj4','obj5'],['obj6','obj7','obj8','obj9']]
    
    it = filter(None, chain.from_iterable(zip_longest(*lists)))
    list(it)
    # ['obj1', 'obj4', 'obj6', 'obj2', 'obj5', 'obj7', 'obj3', 'obj8', 'obj9']
    

    There is also a roundrobin function in the itertools recipes that may be more efficient since it doesn't need to filter the None values, but it's not quite as succinct. You would call that with roundrobin(*lists) to get the desired result.