Search code examples
pythonlist-comprehension

python list comprehension double for


vec = [[1,2,3], [4,5,6], [7,8,9]]
print [num for elem in vec for num in elem]      <----- this

>>> [1, 2, 3, 4, 5, 6, 7, 8, 9]

This is tricking me out.
I understand elem is the lists inside of the list from for elem in vic
I don't quite understand the usage of num and for num in elem in the beginning and the end.

How does python interpret this?
What's the order it looks at?


Solution

  • Lets break it down.

    A simple list-comprehension:

    [x for x in collection]
    

    This is easy to understand if we break it into parts: [A for B in C]

    • A is the item that will be in the resulting list
    • B is each item in the collection C
    • C is the collection itself.

    In this way, one could write:

    [x.lower() for x in words]
    

    In order to convert all words in a list to lowercase.


    It is when we complicate this with another list like so:

    [x for y in collection for x in y] # [A for B in C for D in E]
    

    Here, something special happens. We want our final list to include A items, and A items are found inside B items, so we have to tell the list-comprehension that.

    • A is the item that will be in the resulting list
    • B is each item in the collection C
    • C is the collection itself
    • D is each item in the collection E (in this case, also A)
    • E is another collection (in this case, B)

    This logic is similar to the normal for loop:

    for y in collection:     #      for B in C:
        for x in y:          #          for D in E: (in this case: for A in B)
            # receive x      #              # receive A
    

    To expand on this, and give a great example + explanation, imagine that there is a train.

    The train engine (the front) is always going to be there (the result of the list-comprehension)

    Then, there are any number of train cars, each train car is in the form: for x in y

    A list comprehension could look like this:

    [z for b in a for c in b for d in c ... for z in y]
    

    Which would be like having this regular for-loop:

    for b in a:
        for c in b:
            for d in c:
                ...
                    for z in y:
                        # have z
    

    In other words, instead of going down a line and indenting, in a list-comprehension you just add the next loop on to the end.

    To go back to the train analogy:

    Engine - Car - Car - Car ... Tail

    What is the tail? The tail is a special thing in list-comprehensions. You don't need one, but if you have a tail, the tail is a condition, look at this example:

    [line for line in file if not line.startswith('#')] 
    

    This would give you every line in a file as long as the line didn't start with a hash character (#), others are just skipped.

    The trick to using the "tail" of the train is that it is checked for True/False at the same time as you have your final 'Engine' or 'result' from all the loops, the above example in a regular for-loop would look like this:

    for line in file:
        if not line.startswith('#'):
            # have line
    

    please note: Though in my analogy of a train there is only a 'tail' at the end of the train, the condition or 'tail' can be after every 'car' or loop...

    for example:

    >>> z = [[1,2,3,4],[5,6,7,8],[9,10,11,12]]
    >>> [x for y in z if sum(y)>10 for x in y if x < 10]
    [5, 6, 7, 8, 9]
    

    In regular for-loop:

    >>> for y in z:
        if sum(y)>10:
            for x in y:
                if x < 10:
                    print x
                    
    5
    6
    7
    8
    9