Python List Comprehension: With Examples

Working with arrays can be a bit of a pain in most languages.

If you want to make changes to an array you need to use a FOR loop, a bit like this one:

fruits = ["apple", "banana", "grape", "mango", "pear", "raspberry"]
new_fruits = []

for x in fruits:
  new_fruits.append(x + "s")

print(new_fruits)

All we wanted to do was add an “s” to the end of each of the words in the array but to do that we had to create a new array and then a loop to add the letter to each of the words.

Replacing FOR Loops with List Comprehension

Thanks to the power of List Comprehension in Python, there is a much better way of doing this, which will save you a lot of time and code.

fruits = ["apple", "banana", "grape", "mango", "pear", "raspberry"]
fruits = [fruit + "s" for fruit in fruits]
print(fruits)

The second line contains the Python List Expression which is doing exactly the same thing as the previous example.

We no longer need a new array as we can just overwrite the previous array.

If you want to create your own Python List Expressions then you can use the following formula.

new_list = [expression for item in array if condition == True]

Using Conditionals in List Comprehension

You can see in the formula above that we can also use conditionals in our expressions.

Let’s say we are upset by the fact that “raspberrys” isn’t actually a word. We can add a conditional to a statement that will filter this word out.

fruits = ["apple", "banana", "grape", "mango", "pear", "raspberry"]
fruits = [fruit + "s" for fruit in fruits if fruit != "raspberry"]
print(fruits)

We can also include conditionals at the start if we want to replace certain entries:

fruits = ["apple", "banana", "grape", "mango", "pear", "raspberry"]
fruits = ["raspberries" if fruit == "raspberry" else fruit + "s" for fruit in fruits]
print(fruits)

String Manipulation

Given that strings are nothing more than character arrays we can also use List Comprehension to manipulate text.

Say for example we had the following piece of text and we wanted to add the spaces back in:

ILikeSpellingFruitPluralsIncorrectly

Normally we would need to loop through each letter looking to see if it is upper case or not and add in a space.

text = "ILikeSpellingFruitPluralsIncorrectly"

text = "".join(
    [ i if i.islower() else " " + i for i in text ]
)[1:]
print(text)

Which prints out: I Like Spelling Fruit Plurals Incorrectly

Final Thoughts

List Comprehension is a powerful way of doing simple FOR loops and array manipulations without needing to write a lot of code.

As Uncle Ben said to Peter Parker, “With great power comes great responsibility”. List Comprehension is great for simple expressions to make your code cleaner, but take caution if you are writing complicated loops in this way.

Even if it takes a few more lines of code, it is always better that your code is easier to understand than concise.