List Comprehension in Python

Introduction

In the previous post, we learned all about the basics of Python Lists. As a continuation to that, today we will explore Python List Comprehension. As a beginner learning Python, I was always writing multiple lines of code just to perform simple tasks like creating, updating and modifying lists. Then, I came across List Comprehension that allowed me to perform tasks that required lines of code, in just one clean line.

Even in the real world, whether it is data cleaning or automating tasks, List Comprehension will save a lot of time and effort.

In this blog, I’ll take about List Comprehension, compare it with traditional loops and walk you through some easy-to-understand examples.

What is List Comprehension?

List comprehension is a concise way of creating a new list by applying an expression to each item in an existing iterable. It provides a cleaner and more readable code in comparison to the traditional loop technique.

Basic Syntax:

new_list = [expression for item in iterable]

where, expression = what is to be done to each item

item = each element in iterable

iterable = a list, string or range

Comparing it with Loops:

Let’s suppose you have a list and you want to create a new list containing the cubes of the elements from the original list.

We will first solve this problem using traditional loops and then see how we can achieve the same result using list comprehension in a more simpler, cleaner and easier way.

• Using Loops:

l=[1,2,3,4]
cube=[]
for i in l:
    cube.append(i**3)
print(cube)

#Output
#[1, 8, 27, 64]

• Using List Comprehension:

l=[1,2,3,4]
cube=[i**3 for i in l]
print(cube)

#Output
#[1, 8, 27, 64]

Both methods give us the same output, but using list comprehension:

1. Makes the code easier to read and maintain

2. Requires less code

3. Is often faster in performance

Basic Examples

• Generate a multiplication table of 5

table=[5*x for x in range(1,11)]
print(table)

#Output
#[5, 10, 15, 20, 25, 30, 35, 40, 45, 50]

• Given a list of words, create a list of their first letters.

fruit=['apple','banana','guava','pineapple']
first_letters=[x[0] for x in fruit]

print(first_letters)

#Output
# ['a', 'b', 'g', 'p']

• Convert all words in a list to uppercase.

words=['hello','basket','fruits','python']
uppercase=[x.upper() for x in words]

print(uppercase)

#Output
# ['HELLO', 'BASKET', 'FRUITS', 'PYTHON']

• Given a list of words, repeat each word 3 times in the list

word=['hi','bye','night']
triple=[x*3 for x in word]

print(triple)

#Output
#['hihihi', 'byebyebye', 'nightnightnight']

Conditionals in List Comprehension

So far, we have just used for-statements in our list comprehension. But did you know that we can add a conditional to this, to filter or modify data based on some criteria?

1. Using if:

It can be used to include only those items that meet a certain criteria.

Syntax: new_list = [expression for item in iterable if condition]

Example: To filter out even numbers form a given list.

numbers=[2,5,6,9,0,11,15]
even=[x for x in numbers if x%2==0]

print(even)

#Output
#[2, 6, 0]

2. Using if…else:

It can be used to perform different actions on the items based on the condition.

Syntax: new_list = [expression_if_true if condition else expression_if_false for item in iterable]

Example: To label the numbers of a list as Even or Odd.

numbers=[2,5,6,9,0,11,15]
even_or_odd=['Even' if x%2==0 else 'Odd' for x in numbers ]

print(even_or_odd)

#Output
#['Even', 'Odd', 'Even', 'Odd', 'Even', 'Odd', 'Odd']

3. Using Nested Conditionals:

Sometimes, you might need to apply conditions to include or edit the items from a list. This is where the nested conditional in list comprehension comes in use.

• Using multiple if conditions:

Example: From a given list, filter out those elements which is divisible by 5 as well 10.

l=list(range(1,21))
l1=[x for x in l if x%5==0 if x%10==0]

print(l1)

#Output
#[10, 20]

Note: Two if statements vs if.. and..

In list comprehension, writing:

[x for x in item if condition1 if condition2]

is functionally same as writing,

[x for x in item if condition1 and condition2]

Which one should you use?

• Use if.. and … for better readability and cleaner code.

• Use two if statements when you need to break complex logic into simpler parts.

The above example using if.. and..:

l=list(range(1,21))
l1=[x for x in l if x%5==0 and x%10==0]

print(l1)

#Output
#[10, 20]

Tip: Use if.. and.. , as it is easier to understand and widely used.

• Using nested if else condition:

Example: Grading students based in their marks.

marks=[12,80,67,29,97,48]
grade=['A' if x >= 80 else 'B' if x >= 60 else 'C' for x in marks]

print(grade)

#Output
#['C', 'A', 'B', 'C', 'A', 'C']

Nested List Comprehension

Nested list comprehension means using one list comprehension inside of another. It helps us to work with 2D and 3D lists or to perform some complex tasks in a compact and cleaner way.

Examples:

• Flattening a 2D list.

L=[[1,2],[3,4],[5,6]]
L1=[x for row in L for x in row]

print(L1)

#Output
#[1, 2, 3, 4, 5, 6]

Explanation: First the outer loop goes through each row in list L and then the inner loop goes through the items inside each row. Then, each x is collected and added to the L1 list.

• Creating a 3D list

matrix=[[[i*j*k for i in range(1,3)] for j in range(1,3)] for k in range(1,3)]
print(matrix)

#Output
#[[[1, 2], [2, 4]], [[2, 4], [4, 8]]]

Explanation: First the outer loop (k) runs, that controls how many 2D lists are created.

Then, for each layer, the middle loop (j) runs, that controls the number of rows in the each 2D list.

At last, the innermost loop (i) runs, that generates the actual value inside each row. Each value is created using the expression i * j * k.

• Cartesian product

l1=[1,2]
l2=[5,6,7]

car_pro=[(i,j) for i in l1 for j in l2]
print(car_pro)

#Output
#[(1, 5), (1, 6), (1, 7), (2, 5), (2, 6), (2, 7)]

Explanation: First, the outer-loop (i) runs over each element of l1.

Then, for each value of i, the inner-loop (j) runs over each element of l2.

This will create all the possible ordered pair (i,j), resulting the Cartesian product of the list l1 and l2.

Common Mistakes I Made (So You Don’t Have To)

While learning List Comprehension, I stumbled upon a few problems that confused me a lot. Here are the list of those mistakes, so you don’t make the same ones:

1. Mixing up the order of the loops in nested loop comprehension:

One of my first mistakes was i mixed up the order of the loop in the nested loop comprehension which gave me completely different results. When i first tried using nested list comprehension, i thought the first loop written is always the outer-most loop but i completely overlooked the brackets.

Let me show you what happened exactly.

What I expected: A Cartesian product of the two lists l1 and l2.

What I got: A nested list where each inner list contained tuples with all elements of l1 for a single j from l2.

#My wrong attempt

l1=[1,2]
l2=[5,6,7]

car_pro=[[(i,j) for i in l1] for j in l2] 
print(car_pro)

#Output
#[[(1, 5), (2, 5)], [(1, 6), (2, 6)], [(1, 7), (2, 7)]]

#The correct way
l1=[1,2]
l2=[5,6,7]

car_pro=[(i,j) for i in l1 for j in l2]
print(car_pro)

#Output
#[(1, 5), (1, 6), (1, 7), (2, 5), (2, 6), (2, 7)]

Explanation:

In the wrong version, i mistakenly kept the outer loop in a list that caused a nested list structure.

In the correct version, the loops are written in a proper left to right manner.

What I’ve Learned?

• When there is no brackets, the first loop written (from the left) is the outer-most loop.

• When is brackets i.e. for nested list structure, the brackets determine the order of execution of the loops. The loop inside the bracket is the inner-most loop and the one outside the bracket is the outer-most one.

• Wrapping the comprehension inside another list creates a nested list.

2. Forgetting Brackets:

At first, I would always forget to wrap the list comprehension in brackets.

 #The wrong way
car_pro=(i,j) for i in l1 for j in l2

#The correct way
car_pro=[(i,j) for i in l1 for j in l2]

3. Trying to use .append() inside the list comprehension:

I thought i could use .append() in the list comprehension. But comprehensions returns the new list automatically so, .append() doesn’t work.

#The wrong way
l=[l.append(x) for x in range(1,4)] 
print(l)

#This throws an error as 'l' is not defined yet and we are trying to use l.append(x)

#The correct way
l=[x for x in range(1,4)]
print(l)

#Output
#[1, 2, 3]

4. Over-complicating:

In list comprehension, i tried using too many if.. else conditions in one line making it complex and unreadable. So, you should stick to using the traditional loops if the problem is too complex.

Example: Given a list of fruit names. Create a new list where:

• If the word has more than 5 characters and contains the letter 'e', replace all ‘e‘ with ‘3‘ and convert it to uppercase.

• If the word has less than or equal to 5 characters, reverse the word.

• Otherwise, leave it unchanged.

#Uing list comprehension
l=['apple', 'elephant','cactus', 'dog', 'cheese', 'grape', 'run', 'computer']

new=[x.upper().replace('E','3') if len(x)>5 and 'e' in x else x[::-1] if len(x)<=5 else x for x in l]

print(new)

#Using traditional loop
new_list=[]

for i in l:
    if len(i)>5 and 'e' in i:
        new=i.upper().replace('E','3')
        new_list.append(new)
    elif len(i)<=5:
        new_list.append(i[::-1])
    else:
        new_list.append(i)

print(new_list)

#Output for both
#['elppa', '3L3PHANT', 'cactus', 'god', 'CH33S3', 'eparg', 'nur', 'COMPUT3R']

Since the logic here is complex and included too many conditions, using a traditional loop is preferred, even though writing more lines of code is required.

Using a traditional loop makes the code easier to read, understand and debug compared to a complicated list comprehension.

What’s Next?

Now that we’ve learned all about List Comprehension in Python, it’s time to look behind the scenes.

In the next blog, we will try to understand how lists are stored in memory.

Stay Tuned to learn how understanding memory will help you write efficient code in Python!!