DEV Community

Cover image for Write a List in One Line (List Comprehensions)
Akhilesh
Akhilesh

Posted on

Write a List in One Line (List Comprehensions)

You have been building lists the long way.

numbers = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
squares = []

for n in numbers:
    squares.append(n * n)

print(squares)
Enter fullscreen mode Exit fullscreen mode

Output:

[1, 4, 9, 16, 25, 36, 49, 64, 81, 100]
Enter fullscreen mode Exit fullscreen mode

Four lines. A variable, a loop, an operation, an append. This works perfectly. Nothing wrong with it.

But Python has a shorter way. One line instead of four.

numbers = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
squares = [n * n for n in numbers]

print(squares)
Enter fullscreen mode Exit fullscreen mode

Output:

[1, 4, 9, 16, 25, 36, 49, 64, 81, 100]
Enter fullscreen mode Exit fullscreen mode

Same result. One line. This is a list comprehension.


Reading It Out Loud

The trick to understanding list comprehensions is reading them left to right like a sentence.

[n * n for n in numbers]
Enter fullscreen mode Exit fullscreen mode

"Give me n * n for each n in numbers."

That's literally what it says. The expression comes first. Then the loop. Your brain wants the loop first because that's how you think about it as code. But the comprehension puts the result first, then explains where it comes from.

Take a minute with that. Once it clicks, it stays clicked.


The Basic Shape

[expression for variable in iterable]
Enter fullscreen mode Exit fullscreen mode

expression is what you want in the new list. Can be anything. A calculation, a function call, a transformation.

variable is the loop variable. Same as for n in ....

iterable is what you're looping over. A list, a range, a string, anything.

names = ["alex", "priya", "sam", "jordan"]
upper_names = [name.upper() for name in names]
print(upper_names)
Enter fullscreen mode Exit fullscreen mode

Output:

['ALEX', 'PRIYA', 'SAM', 'JORDAN']
Enter fullscreen mode Exit fullscreen mode
lengths = [len(name) for name in names]
print(lengths)
Enter fullscreen mode Exit fullscreen mode

Output:

[4, 5, 3, 6]
Enter fullscreen mode Exit fullscreen mode
numbers = range(1, 6)
doubled = [n * 2 for n in numbers]
print(doubled)
Enter fullscreen mode Exit fullscreen mode

Output:

[2, 4, 6, 8, 10]
Enter fullscreen mode Exit fullscreen mode

Adding a Condition

You can filter which items make it into the new list.

[expression for variable in iterable if condition]
Enter fullscreen mode Exit fullscreen mode

The if at the end acts as a filter. Only items where the condition is True get included.

numbers = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
evens = [n for n in numbers if n % 2 == 0]
print(evens)
Enter fullscreen mode Exit fullscreen mode

Output:

[2, 4, 6, 8, 10]
Enter fullscreen mode Exit fullscreen mode

Read it: "Give me n for each n in numbers, but only if n is even."

scores = [45, 92, 78, 61, 34, 88, 55, 97]
passing = [s for s in scores if s >= 60]
print(passing)
Enter fullscreen mode Exit fullscreen mode

Output:

[92, 78, 61, 88, 97]
Enter fullscreen mode Exit fullscreen mode

Six lines of loop code, one line of comprehension. Both do exactly the same thing. The comprehension just says it more directly.


Combining Both

Transform and filter at the same time.

numbers = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
squared_evens = [n * n for n in numbers if n % 2 == 0]
print(squared_evens)
Enter fullscreen mode Exit fullscreen mode

Output:

[4, 16, 36, 64, 100]
Enter fullscreen mode Exit fullscreen mode

"Give me n * n for each n in numbers, but only if n is even." Even numbers are 2, 4, 6, 8, 10. Their squares are 4, 16, 36, 64, 100.


With Strings

List comprehensions work on anything iterable. Including strings themselves, since a string is just a sequence of characters.

sentence = "Hello World"
letters = [char for char in sentence if char != " "]
print(letters)
Enter fullscreen mode Exit fullscreen mode

Output:

['H', 'e', 'l', 'l', 'o', 'W', 'o', 'r', 'l', 'd']
Enter fullscreen mode Exit fullscreen mode
words = ["  hello  ", "  world  ", "  python  "]
cleaned = [w.strip() for w in words]
print(cleaned)
Enter fullscreen mode Exit fullscreen mode

Output:

['hello', 'world', 'python']
Enter fullscreen mode Exit fullscreen mode

Cleaning whitespace from a list of strings in one line. You'll do this constantly when processing real data.


Dictionary Comprehensions

Same idea, curly braces instead of square brackets, and you define both the key and the value.

names = ["Alex", "Priya", "Sam"]
name_lengths = {name: len(name) for name in names}
print(name_lengths)
Enter fullscreen mode Exit fullscreen mode

Output:

{'Alex': 4, 'Priya': 5, 'Sam': 3}
Enter fullscreen mode Exit fullscreen mode
scores = {"Alex": 92, "Priya": 58, "Sam": 75, "Jordan": 44}
passing = {name: score for name, score in scores.items() if score >= 60}
print(passing)
Enter fullscreen mode Exit fullscreen mode

Output:

{'Alex': 92, 'Sam': 75}
Enter fullscreen mode Exit fullscreen mode

Build a new dictionary from an existing one, filtered or transformed. Very clean.


When Not to Use Them

Comprehensions are great. They are not always the right choice.

If the logic inside is getting complex, a regular loop is clearer.

# getting hard to read
result = [process(transform(validate(x))) for x in data if check_one(x) and check_two(x)]

# clearer as a loop
result = []
for x in data:
    if check_one(x) and check_two(x):
        validated = validate(x)
        transformed = transform(validated)
        result.append(process(transformed))
Enter fullscreen mode Exit fullscreen mode

The rule is readability. If someone can read the comprehension and understand it in three seconds, use it. If they have to stop and decode it, write a loop.


Try This

Create comprehensions_practice.py.

Start with this data:

temperatures_c = [0, 10, 20, 30, 40, 100]
words = ["Python", "is", "great", "for", "AI", "and", "ML"]
students = [
    {"name": "Alex", "score": 88},
    {"name": "Priya", "score": 52},
    {"name": "Sam", "score": 76},
    {"name": "Jordan", "score": 91},
    {"name": "Lisa", "score": 43}
]
Enter fullscreen mode Exit fullscreen mode

Do all of this using comprehensions only, no regular loops:

Convert every temperature to Fahrenheit and store in a new list.

Build a list of words that are longer than 3 characters.

Build a list of just the names from the students list.

Build a list of just the names of students who passed (score 60 or above).

Build a dictionary mapping each student's name to their score.


What's Next

List comprehensions make your code shorter. The next post is about lambda functions and map and filter, which are different tools for similar problems. They come from a style of programming called functional programming and they show up constantly in data science code.

Top comments (0)