Python - List Comprehension

I am a Software and Data Engineer with great passion for software development and making data science delightful.
In this article, we are going to learn about list comprehension in python using comprehensive examples. You will learn about:
- List Comprehensions and how they are different from For Loop.
- Replacing For Loops with list comprehensions.
- Adding conditions to list comprehensions.
- When NOT to use List Comprehension.
What is List Comprehension?
List comprehension is a programming syntax in python used for creating new lists from other lists and iterables like tuples, strings, arrays, etc. It consists of square brackets containing the expression, which is executed for each element along with the for loop to iterate over each element.
Benefits of List Comprehension
- More concise way to create lists in python.
- More time and space-efficient than loops.
- More Pythonic way to create, modify, and filter lists.
Writing List Comprehension
List Comprehension syntax:
newlist = [expression for item in iterable]
For example:
newlist = [i**2 for i in range(5)] # creates a list of the squares of 0 to 4
print(newlist)
Output:
[0, 1, 4, 9, 16]
Every list comprehension in Python includes three elements:
- Expression: is the member itself, a call to a method, or any other valid expression that returns a value. The expression in the above example is i**2.
- Item: this is the object or value in the list or iterable. The item in the above example is i.
- Iterable: this is a list, set, sequence, generator, or any other object that can return its elements one at a time. The iterable in the example above is 5.
List Comprehension vs For Loop:
There are various ways of iterating through a list, however, For loop is the most commonly used approach. Check out the example below.
oldlist = ['apple', 'ball', 'cat', 'dog', 'egg']
newlist = []
# Iterating using for loop
for word in oldlist:
if "a" in word:
newlist.append(word)
print(newlist)
Output:
['apple', 'ball', 'cat']
We can simply rewrite the above code in List Comprehension as shown below:
oldlist = ['apple', 'ball', 'cat', 'dog', 'egg']
newlist = [word for word in fruits if "a" in word]
print(newlist)
Output:
['apple', 'ball', 'cat']
Time comparison between For Loop and List Comprehension:
By importing the time module, we can compare the execution time of a For loop and List Comprehension as shown below.
# Import time module
import time
# Function to implement for loop
def for_loop(n):
result = []
for i in range(n):
result.append(i**2)
return result
# Function to implement list comprehension
def list_comprehension(n):
return [i**2 for i in range(n)]
# Calculate time taken by for_loop()
begin = time.time()
for_loop(10**6)
end = time.time()
# Display time taken by for_loop()
print('Time taken for_loop:',round(end-begin,2))
# Calculate time takens by list_comprehension()
begin = time.time()
list_comprehension(10**6)
end = time.time()
# Display time taken by for_loop()
print('Time taken for list_comprehension:',round(end-begin,2))
Output:
Time taken for_loop: 0.56 Time taken for list_comprehension: 0.47
Using if...else with list comprehension
We can also add 'if...else' condition to list comprehension
num = ["Even" if i%2==0 else "Odd" for i in range(5)]
print(num)
Output:
['Even', 'Odd', 'Even', 'Odd', 'Even']
Here, list comprehension will check the 5 numbers from 0 to 4. If i is divisible by 2, then 'Even' is appended to the num list. If not, 'Odd' is appended.
Nested List comprehension
nested_list = [[1,2,3],[4,5,6],[7,8,9]]
flat_list = []
for list in nested_list:
for item in list:
flat_list.append(item)
print(flat_list)
To do this with a list comprehension, we can write:
nested_list = [[1,2,3],[4,5,6],[7,8,9]]
flat_list = [i for j in nested_list for i in j]
print(flat_list)
Both of these return:
[1, 2, 3, 4, 5, 6, 7, 8, 9]
When Not to Use List Comprehensions
While in most cases, list comprehensions tend to be more readable than other methods, they can also become quite complex. When this occurs, it may be better to use another more readable method.
Even though every list comprehension can be written as a for loop, it is important to note that not every for-loop can be a list comprehension.
Conclusion
- List comprehension is an elegant way to define and create lists based on existing lists and iterables like tuples, strings, arrays, etc.
- Generally, list comprehension is more lightweight and compact than standard list formation functions and loops.
- We should not write long codes for list comprehensions in order to ensure user-friendly code.
- Remember, every list comprehension can be rewritten in for loop, but every for loop can’t be rewritten in the form of list comprehension.


