Python is famous for its elegant coding style that is concise, easy to write and almost as easy to read as plain English. List comprehension in python is one of the distinctive feature which provides a powerful functionality within a single line of code. It’s one of the advanced topic of Python that is very easy to learn, implement and practically useful. It makes your code concise, readable and boost speed in most cases.
What is List Comprehension?
List comprehension is a tool for creating a list from other list or sequence. Elements can be conditionally included in the new list and each element can be transformed as needed. It is much more concise and readable alternatives to creating a list in python.
Let’s see an example where you want to create a list containing a square of numbers from 1 to 10.
square = [] # empty list for storing square of number
for i in range(1, 11):
square.append(i*i)
List comprehension is a syntactic sugar that replaces the above code into following a single line of code
square = [ i*i for i in range(1,11)]
The above code is just the tip of an iceberg. You can do a lot more using list comprehension.
Why use list comprehension in Python
- Concise: Multiple line of code can be replaced with one liner list comprehension
- Boost Speed: Not only, it is concise; it is faster in most cases.
- Readability: Improves overall readability of code
Syntax
#Above Example is equivalent to
cube = []
for i in [1, 2, 3, 4]:
if i > 2:
cube.append(i**3)
# cube >>> [27 64]
In the example,
- Output: It is the value that is stored in the list. It is an expression, or function of variable used in for loop. The variable of each iteration can be manipulated if required in the output.
- Variable: For loop iterates through each item of list/sequence and stores it in variable.
- Condition (Optional): Although optional, you can filter the outcome of for loop using condition. Variable that satisfies the condition is only available for output.
Manipulating output of list expression:
Following code generates a list of number from 1 to N. This is the base case that we will build upon and add complexity
N = 5
numbers = [i for i in range(1, N+1)] # Example of simple list comprehension
print(numbers)
# Output
>>> [1, 2, 3, 4, 5]
# Above code is equivalent to
N = 5
numbers = []
for i in range(1, N+1):
numbers.append(i)
print(numbers)
# Output
>>> [1, 2, 3, 4, 5]
To get a list of square, just replace i in output with i*i
N = 5
numbers = [i*i for i in range(1, N+1)] # Example of simple list comprehension
print(numbers)
# Output >>> [1, 4, 9, 16, 25]
# Above code is equivalent to
N = 5
numbers = []
for i in range(1, N+1):
numbers.append(i*i)
print(numbers)
# Output >>> [1, 4, 9, 16, 25]
Instead of manipulating variable in output, any functions or expression can be used. In the following example shows how to use function.
N = 5
def f(x):
return (x*x + 2*x + 1)
numbers = [f(i) for i in range(1, N+1)]
print(numbers)
# Output >>> [4, 9, 16, 25, 36]
# Above code is equivalent to
N = 5
numbers = []
def f(x):
return (x*x + 2*x + 1)
for i in range(1, N+1):
numbers.append(f(i))
print(numbers)
# Output >>> [4, 9, 16, 25, 36]
Creating Nested list with Multiple output: Returning multiple output as a list or tuple is also possible with list comprehension. This is very useful for creating function values for a range of input.
N = 3
def f(x):
return (x*x + 2*x + 1)
numbers = [[x,f(x)] for x in range(1, N+1)]
print(numbers)
# Output >>> [[1, 4], [2, 9], [3, 16]]
# Above code is equivalent to
N = 3
numbers = []
def f(x):
return (x*x + 2*x + 1)
for x in range(1, N+1):
numbers.append([x, f(x)])
print(numbers)
# Output >>> [[1, 4], [2, 9], [3, 16]]
Iterating over sequence
Next, in list comprehension, you can iterate over any sequence or iterable like list, generator, dictionary, tuple or string.
Example of of iterating over list of numbers and convert each element of a list into upper case.
numbers = ['One', 'Two', 'Three', 'Four', 'Five']
uppercase_numbers = [num.upper() for num in numbers]
print(uppercase_numbers)
# Output >>> ['ONE', 'TWO', 'THREE', 'FOUR', 'FIVE']
Iteraing over dictonary
countries = {'India': 'New Delhi', 'US': 'Washington DC', 'UK': 'London'}
capital = [countries[key] for key in countries]
print(capital)
# Output >>> ['New Delhi', 'Washington DC', 'London'
Nested Loop in List Comprehension
Till now we have used only one for loop. However, it is possible to use nested for loop too in list expression. In the following example, the first for loop extracts the rows of nested list, while the second for loop extracts numbers from each row which is being stored in the list.
vec = [[1,2,3], [4,5,6], [7,8,9]]
v = [num for row in vec for num in row]
print(v)
# Output >> [1, 2, 3, 4, 5, 6, 7, 8, 9]
# Above code is equivalent to
vec = [[1,2,3], [4,5,6], [7,8,9]]
v = []
for rowin vec:
for num in row:
v.append(num)
print(v)
# Output >> [1, 2, 3, 4, 5, 6, 7, 8, 9]
Using if condition in list comprehension
Although optional, you can create a list of elements that follows particular condition. Following example filters out even numbers or a positive number from another list.
# Filter out Even numbers as list
sample1 = [2, 4, 7, 6, 9, 15, 20, 4, 7, 8]
even = [num for num in sample1 if num % 2 == 0]
print(even)
# Output >>> [2, 4, 6, 20, 4, 8]
# Filter out positive numbers from list
sample2 = [1, -2, 3, -8. -4, 5, 6, -8, -10, 2]
positive = [num for num in sample2 if num >=0]
print(positive)
# Output >>> [1, 3, 5, 6, 2]
Using if else condition in list comprehension
You can also use ternary operator of python for conditionally filtering data using if else condition. Following code show how it can create a list containing True if number is Even or False if number is not even based on numbers from another list.
# Example of if else condition in list comprehension in Python
numbers = [1,4,5,9,2]
is_even = [True if i % 2 == 0 else False for i in numbers]
print(is_even)
# Output >>> [False, True, False, False, True]
Finding Unique value of list using List comprehension
Finding the unique values of the list using list comprehension with condition
# Finding unique element of list
lst = [1, 7, 3, 4, 1, 5, 6, 3, 0, 2]
unique = []
[unique.append(num) for num in lst if num not in unique]
print(unique)
# Output >>> [1, 7, 3, 4, 5, 6, 0, 2]
Summary
List comprehension in python provides a crisp way to create lists from the other list of iterable like tuple, generator (like range), dictionary and string.
Key points to note about list comprehension are
- It is an elegant way to create lists based on existing lists.
- It is concise, and generally faster than normal functions, and loops for creating list.
- However, avoid writing long and complicated list comprehensions to ensure code readability,.
- Remember, every list comprehension can be rewritten in for loop, but every for loop can’t be rewritten in the form of list comprehension.