Python List Comprehension

If you are learning how to code in python, it is inevitable that you come across lists. List is one of the most versatile, fundamental, data structures in python and used in many applications of data manipulation, machine learning and AI.

Lists involves putting values in between a square bracket - [ ], separated by commas. Example,

planets = ['mars', 'jupiter', 'earth', 'pluto', 'uranus', 'saturn', 'mercury']

Lists are quite easy to understand and manipulate. What tends to be confusing is list comprehension. List comprehension is a ‘pythonic’ way of creating shorter, concise lists. They work just like lists and can be used for iterations and for short conditional expressions. They are a worthy substitute for map(), reduce() and filter() functions.
For the purpose of this article, we will look at list comprehensions as opposed to lists, when and how we can use them.

Syntax

[expression for item in list if conditional]

The syntax for is the first thing that grabs your attention about list comprehension. It seems to go against the normal order for writing python code because the expression comes before the loop. The way list comprehension is written gives it the faster advantage over list. Here, python allocates the list’s memory first instead of having to resize the list on runtime.

From the syntax, all list comprehension has a for loop in it. They are basically iterables to produce another list.

List comprehension and for loop

List comprehensions provide a concise way to create lists. Common applications are to make new lists where each element is the result of some operations applied to each member of another sequence or iterable, or to create a subsequence of those elements that satisfy a certain condition.

Python Docs for Python 3

Although all list comprehension can be written as a for loop, not every for loop can be rewritten as a list comprehension. This is a limitation for list comprehension. Complex for loop statements can take its toll on it.
Here is an example that highlights the difference in using for in list and list comprehension,

names = []
  for name in "Beyonce Knowles":
    names.append(name)
    print(names) 
#['B', 'e', 'y', 'o', 'n', 'c', 'e', ' ', 'K', 'n', 'o', 'w', 'l', 'e', 's']

The above is an example of the typical example of how for-loops are used to iterate strings. It can be rewritten using list comprehension like this,

her_name = [names for names in 'Beyonce Knowles']
print(her_name) 

List comprehension with in-built functions

List comprehensions can also be used for mathematical calculations and other iterations. It works perfectly with functions like range(), enumerate(), zip() and len().
Here is a code snippet to get all the even numbers in a range of 0 to 10.

add = [] 
  for x in range(0, 10, 2): 
    add.append(x) 
...
print(add)
# [2, 4, 6, 8]

Rewritten in list comprehension, it becomes

add = [x for x in range(0,10,2)] 
print(add)
# [2, 4, 6, 8]

The len() function is not iterable as a rule. By wrapping it in a range(), we can get the length of a given string, tuple, or list in an expression. The output, though, is not a single digit but a list of numbers that match up to the number of letters in the string. For the word ‘Beyonce’, instead of getting a 6, when we use len wrapped in range(), we get [1,2,3,4,5,6].

keys = [m for m in range(len("Beyonce"))] 
print(keys) 
#[0, 1, 2, 3, 4, 5, 6]

To use enumerate() in a list comprehension, wrap the variables in a bracket so python can identify it. This is also the syntax to use for zip().

mynewlist = ['boat', 'ship', 'yacht'] 
list = [(e, f) for e, f in enumerate(mynewlist)] 
print(list) 
#[(0, 'boat'), (1, 'ship'), (2, 'yacht')] 

Slicing with list comprehension

Slicing is an indexing operation where indexes of mutable sequences are modified or moved. With slicing, sub-strings or subsets can be obtained from a given value. We can use slicing with list comprehension like this to get the word “Super Star”

word = "I am a SuperStar" 
words = [word[7:12] + " " + word[12:] for i in range(len(word))] 
print(words<i>)  
#Super  Star

With indexing and list comprehension, we can get the first letters or an acronym of a list of strings.

who = [acry[0] for acry in ("World", "Health", "Organization") ] 
print (who) 
#['W', 'H', 'O']

Conditional Statements with List comprehensions

List comprehension is powerful because it utilizes and accommodates the strengths of conditional statements and iterations at the same time. The standard syntax for list comprehension also includes if conditionals

[expression for item in list if conditional]

Using IF with List Comprehension

If statement is used when there is one condition to be met. We will look at the longer syntax and the list comprehension syntax. 

for x in range(20): 
  if x%2==0: 
    print(x)

With list comprehension, this becomes

number = [a for a in range(20) if a%2 == 0] 
print(number)

The output of the above code displays a range of even numbers from two to twenty.

Using IF..ELSE with List Comprehension

IF..Else statements present two conditions to the system. If you can’t do A, then you should do B. List comprehension makes allowance for this in its syntax and here is an example. 

obj = ["Even" if i % 2 == 0 else "Odd" for i in range(10)] 
print(obj) 
#['Even', 'Odd', 'Even', 'Odd', 'Even', 'Odd', 'Even', 'Odd', 'Even', 'Odd']

In just one line, in a range of numbers from one to ten, we have been able to identify the even and odd numbers using list comprehension.

In this example, the if conditional is no longer at the end of the statement. As the conditional gets more complex, the syntax becomes more flexible, allowing for readability and efficiency.

Using ELIF with List Comprehension

ELIF is used when there are usually two or more conditions to be met. The conditions could go on and on depending on the number of conditions to be met. Fizzbuzz, the popular programming task is the perfect example of an ELIF problem that can be solved with list comprehension.

In Fizzbuzz, if a number is divisible by 3, it prints fizz. If it is divisible by 5, it prints buzz. If is divisible by both 3 or 5, it prints Fizzbuzz, and if a number doesn’t fit any of these criteria, it just prints the number.

Using just for loop with conditionals, we will write lines and lines of iterations like this,

def fizz(fizzbuzz): 
for fizzbuzz in range(fizzbuzz): 
  if fizzbuzz % 3 == 0 and fizzbuzz % 5 == 0: 
    print("fizzbuzz") 
      continue 
  elif fizzbuzz % 3 == 0: 
    print("fizz") 
      continue 
  elif fizzbuzz % 5 == 0: 
    print("buzz") 
      continue 
  else: 
    print(fizzbuzz)     
fizz(16) 
#['fizzbuzz', 1, 2, 'fizz', 4, 'buzz', 'fizz', 7, 8, 'fizz', 'buzz', 11, 'fizz', 13, 14, 'fizzbuzz']

But this code can be made shorter with list comprehension, like this

fizzbuzz = ['fizzbuzz' if (num%3==0 and num%5==0) else 'fizz' if (num%3 ==0) else 'buzz' if (num%5 ==0) else num for num in range(16) ] 
print(fizzbuzz)
#['fizzbuzz', 1, 2, 'fizz', 4, 'buzz', 'fizz', 7, 8, 'fizz', 'buzz', 11, 'fizz', 13, 14, 'fizzbuzz']

Note: When writing elif statements conditions in list comprehension, the last condition should be written first, else it will be overlooked. For example, if the fizzbuzz code is written like this, 

fizzbuzz = ['fizz' if num%3==0 else 'buzz' if num%5 ==0 else 'fizzbuzz' if num%3==0 and num%5==0 else num for num in range(16)] 
print(fizzbuzz)

The output will be

# ['fizz', 1, 2, 'fizz', 4, 'buzz', 'fizz', 7, 8, 'fizz', 'buzz', 11, 'fizz', 13, 14, 'fizz']

0 and 15 should display ‘fizzbuzz’, but they have been overridden by the condition for num % 3 == 0.

Using Nested if statements with List Comprehension

Nested if statements are if statements in another if statement.

adj = ['red', 'big', 'tasty'] 
fruits = ['apple', 'cherry', 'banana'] 
  for x in adj: 
    for y in fruits: 
      print(x, y)

In list comprehension, nested conditional like this can be rewritten in one line as, 

adj = ['red', 'big', 'tasty'] 
fruits = ['apple', 'cherry', 'banana']
newadj = [(y, x) for y in adj for x in fruits]
print(newadj)
# [('red', 'apple'), ('red', 'cherry'), ('red', 'banana'), ('big', 'apple'), ('big', 'cherry'), ('big', 'banana'), ('tasty', 'apple'), ('tasty', 'cherry'), ('tasty', 'banana')] 

List Comprehension VS Map(), Reduce() and Filter()

Map(), Filter() and Reduce() functions are widely popular in python and are used extensively. Code snippets shown here can also be rewritten using map(), reduce(), filter(). For example, we can use map to multiply the numbers in the list by 10

nos = [2,3,4,5,6,7,8,9] 
newnos = map(lambda x: x * 10, my_list) 
print(list(newnos))
# [20, 30, 40, 50, 60, 70, 80, 90] 

And we can use filter to display only even numbers from a list

nos = [1,2,3,4,5,6,7,8,9,10]  
is_even = lambda x: x % 2 == 0 
# using filter  
new_even_nos = list(filter(is_even, lis1))  
print(new_even_nos)

List comprehension has advantage over the aforementioned functions. It is faster, concise, more readable, versatile, and does not have to be attached to lambda functions. It is recommended for a more pythonic coding practice in the Python Style Guidelines. 

Things to note about list comprehension

  • The output is always a list
  • List comprehensions are usually one-liners, meaning they are usually written in one line as a shorter, more concise way of re-writing lists.
  • They are not shortcuts or a back-handed way of writing code. They are actually pythonic recommended in the Python Style Guidelines.
  • It is written in square brackets - [ ].
  • An overuse of list comprehension can make a code unreadable, which undermines one of the biggest strengths of list comprehensions. It could also raise efficiency and maintainability issues. 
  • They are faster than for loops.
  • List comprehensions can be used with if…else statements, loops, nested loops and unconditional statements. 

When you should use list comprehensions

  • For simple, concise simple logic. Conditionals that are complex will lead to problems if used with list comprehensions.
  • For simple datasets. Using a generator works best for large datasets.
  • When you don’t have to declare a function. Maps are better used with functions.

Conclusion

Understanding list comprehension will help you write more pythonic code, but list comprehensions also tend to be over and inappropriately used. That is one thing to watch out for. The speed, readability and lightness of list comprehension makes it a notable alternative for lambda functions, map(), reduce() and filter(). This article covered list comprehensions, especially it’s usage with conditional statements.

Author

David Dorr, Head of eCommerce

David is the Head of e-Commerce at CodeCoda where he is responsible to lead several teams of eCommerce specialists. In his previous role as a data scientist for London Metropolitan Police, he was developing deep learning NLP algorithms as part of the Crime Prediction initiative. He then switched over to combine AI with e-Commerce.
He received a B.Sc in Physics from the University of Surrey, Guildford in 1996. With this scientific background, he switched relatively early in his life towards Neural Networks and e-Commerce and has ever since been fascinated with what AI and Machine Learning can do for Online Commerce.