List Comprehension Reference in Python

Recently I took a few moments to play around with some list comprehensions and document a few references that I would have found useful many months ago when I was first learning about and understanding them.  Most recently I still didn’t fully understand the syntax when working with nested loops in a list comprehension, and so I took the time to write out some samples.

The basic pattern is as follows:

Here are some examples written out in both the non-python way across multiple lines and with a one-line list comprehension:

def copy_list(input_list):
    new_list = []
    for item in input_list:
        new_list.append(item)

    new_list = [item for item in input_list]

    return new_list

Starting with the most basic usage, we can copy the items in a list *(without using the copy function).  Here we just use the basic pattern in a list comprehension to iterate over the items in a list.

def copy_items_less_than_number(input_list, threshold_number):
    new_list = []
    for item in input_list:
        if item < threshold_number:
            new_list.append(item)

    new_list = [item for item in input_list if item < threshold_number]

    return new_list

Now we can apply the optional predicate wherein we supply a condition to determine whether or not an item should be included within the list itself.

def modify_list_less_than_number(input_list, threshold_number, new_number):
    new_list = []
    for item in input_list:
        if item < threshold_number:
            new_list.append(new_number)
        else:
            new_list.append(item)

    new_list = [new_number if item < threshold_number else item for item in input_list]

    return new_list

This is not the same as adding conditional statements to determine the value placed in the list.  In this example, the list size does not change, but the values added to the list are dynamic.

def every_other_item(input_list):
    new_list = []
    for index, item in enumerate(input_list):
        if index % 2 == 0:
            new_list.append(item)
    new_list = [item for index, item in enumerate(input_list) if index % 2 == 0]

    return new_list

We can also include additional variables that are in scope within the list comprehension.  If we enumerate a list, we can get each value’s corresponding index as well, and apply logic to use that index.

def copy_dictionary(input_dictionary):
    new_dict = {}
    for key, value in input_dictionary.items():
        new_dict[key] = value

    new_dict = {k: v for k, v in input_dictionary.items()}

    return new_dict

We can also build dictionaries with list comprehensions in Python 2.7 and above.

def invert_dictionary(input_dictionary):
    new_dict = {}
    for key, value in input_dictionary.items():
        new_dict[value] = key

    new_dict = {v: k for k, v in input_dictionary.items()}

    return new_dict

How the variables are used inside the list comprehension is up to you…

def flatten_2d_matrix(two_d_list):
    new_list = []
    for sublist in two_d_list:
        for item in sublist:
            new_list.append(item)

    new_list = [item for sublist in two_d_list for item in sublist]

    return new_list

The above example is the reason I took the time to explore some examples. The syntax didn’t quite make sense to me, and every time I wanted to flatten a list I went and referenced a StackOverflow Question. However, it begins to make sense if we expand this out beyond two dimensions:

def flatten_3d_matrix(three_d_matrix):
    new_list = []
    for two_d_matrix in three_d_matrix:
        for one_d_matrix in two_d_matrix:
            for item in one_d_matrix:
                new_list.append(item)

    new_list = [item for two_d_matrix in three_d_matrix for one_d_matrix in two_d_matrix for item in one_d_matrix]

    return new_list

You’ll notice that the pattern here is to simply take all of your nested loops and just place them one after another without new spacing or line breaks. The rest of the syntax for the list comprehension obviously does not change.