List comprehension behavior with 2 for loop

Viewed 281

Here is the code to emulate python zip function.

def myzip(*args):
    return [tuple(a[i] for a in args) for i in range(len(min(args, key=len)))]

print(myzip([10, 20,30], 'abc'))   

#output: [(10, 'a'), (20, 'b'), (30, 'c')]

If I remove the tuple(), the output will be: [10, 20, 30, 'a', 'b', 'c']

I don't quite understand how list comprehension works when we just add tuple()

So after each 2 loop, it yield a value and automatically add it to internal list before turning into tuple and finally add them to outer list?

For eg:

Loop 1: 10 Loop 2: a Add to [10,a] -> tuple([10,a]) -> (10,a)

Loop 2: 20 Loop 2: b Add to [20,b] -> tuple([20,b]) -> (20,b)

....

3 Answers

If you remove tuple you are no longer making the individual sub lists (or tuples) and you have changed the inner vs outer loops in the comprehension.

This comprehension works (in a similar fashion to tuple(a[i] for a in args)):

[[a[i] for a in args] for i in range(len(min(args, key=len)))]
# [[10, 'a'], [20, 'b'], [30, 'c']] with myzip([10,20,30],'abc')

since the [a[i] for a in args] is executed entirely before the next i; ie, each i'th index in each a is added to a sub list in turn and a sublist of those results is added to the resulting list. Hence [a[i] for a in args] is the inner loop and for i in range(len(min(args, key=len))) is the outer loop.

But this:

 [a[i] for a in args for i in range(len(min(args, key=len)))]
 # [10, 20, 30, 'a', 'b', 'c'] with myzip([10,20,30],'abc')

exhausts the i list before going to the next a in args. The i list is then restarted and explains why you get [10, 20, 30, 'a', 'b', 'c'] instead of [(10, 'a'), (20, 'b'), (30, 'c')] Removing either [] or tuple() has now made for i in range(len(min(args, key=len))) the inner loop and a[i] for a in args has now become the outer loop.


And this:

[(a[i] for a in args) for i in range(len(min(args, key=len)))]

creates a list of generators.

If I remove the tuple(), the output will be: [10, 20, 30, 'a', 'b', 'c']

I assume you meant

>>> def myzip(*args):
...     return [a[i] for a in args for i in range(len(min(args, key=len)))]
...
>>> print(myzip([10, 20,30], 'abc'))
[10, 20, 30, 'a', 'b', 'c']

which is equivalent to

>>> def myzip(*args):
...     results = []
...     for a in args:
...         for i in range(len(min(args, key=len))):
...             results.append(a[i])
...     return results
...
>>> print(myzip([10, 20,30], 'abc'))
[10, 20, 30, 'a', 'b', 'c']

This essentially flattens the nested list in args. See this answer on more about "nested" list comprehension.

This just has to do with where you position the different parts of the list comprehension expression.

With the tuple() call in place, you are calling the tuple function and passing the generator (a[i] for a in args). So, automatically, the portion for a in args will be considered the inner loop. Hence, [(10, 'a'), (20, 'b'), (30, 'c')].

This can be expanded to:

>>> ret = []
>>> for i in range(len(min(args, key=len))):
        sub = []
        for a in args:
            sub.append(a[i])
        ret.append(tuple(sub))

    
>>> ret
[(10, 'a'), (20, 'b'), (30, 'c')]

Without the tuple() call in place, the portion for i in range(len(min(args, key=len))) becomes the inner loop. Hence, [10, 20, 30, 'a', 'b', 'c'].

This can be expanded to:

>>> ret = []
>>> for a in args:
        for i in range(len(min(args, key=len))):
            ret.append(a[i])

        
>>> ret
[10, 20, 30, 'a', 'b', 'c']

If you were two switch the two loops to make this:

[a[i] for i in range(len(min(args, key=len))) for a in args]

Now, the portion for a in args is back to the inner loop. Hence, [10, 'a', 20, 'b', 30, 'c'].

This can be expanded to:

>>> ret = []
>>> for i in range(len(min(args, key=len))):
        for a in args:
            ret.append(a[i])

        
>>> ret
[10, 'a', 20, 'b', 30, 'c']
Related