List Comprehensions for Data Manipulation

From for Loop to List Comprehension

Let's start with a familiar task from the previous lesson: building a new list from an existing one. Imagine you want to create a list containing the squares of numbers from 0 to 9. Using a for loop, you would write something like this:

squares = []
for x in range(10):
    squares.append(x**2)
    
print(squares)



# Output: [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

This pattern is so common that Python provides a special syntax to condense it. A list comprehension achieves the same result in one line:

squares = [x**2 for x in range(10)]

print(squares)



# Output: [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

Notice how the key components of the for loop are still present, just rearranged:

• x**2 is the expression — what to do with each item.
• for x in range(10) is the loop — which iterates over the source sequence.

The following video from Real Python provides an excellent step-by-step breakdown of how to transform a for loop into a list comprehension.

Adding a Condition to Filter Data

List comprehensions become even more powerful when you add conditional logic. Suppose you only want the squares of the odd numbers. With a for loop, you'd add an if statement:

odd_squares = []
for x in range(10):
    if x % 2 != 0:  # Check if x is odd
        odd_squares.append(x**2)

print(odd_squares)



# Output: [1, 9, 25, 49, 81]

In a list comprehension, you can add an if clause at the end to achieve the same filtering effect:

odd_squares = [x**2 for x in range(10) if x % 2 != 0]

print(odd_squares)



# Output: [1, 9, 25, 49, 81]

The general structure is [expression for item in iterable if condition]. This is visually summarized in the following diagram.

The mathspp blog offers a fantastic, in-depth article on this topic. Let's read the sections that cover this core syntax and the addition of filtering.

A Connection to Mathematical Notation

Given your background in mathematics, you'll likely appreciate the deep connection between list comprehensions and set-builder notation. In mathematics, you might define the set of the first ten square numbers like this:

This reads as "S is the set of all numbers $n^2$ such that $n$ is an element of the set of integers from 0 to 9."

Now look at the Python list comprehension again:
squares = [n**2 for n in range(10)]

The parallel is striking. The Python syntax is a direct, executable implementation of the mathematical concept.

• The expression $n^2$ corresponds to n**2.
• The vertical bar | (meaning "such that") and the for...in construct serve the same purpose.
• The source set $n\in \{0,1,\dots ,9\}$ corresponds to n in range(10).

This extends to filtering as well. If we wanted to define only the squares of the odd numbers, we might write:

The Python equivalent mirrors this beautifully:
odd_squares = [n**2 for n in range(10) if n % 2 != 0]

The mathspp article you started reading has a section that explores exactly this connection.

Readability and When to Use Comprehensions

List comprehensions are favored in Python because they are often more concise and readable than the equivalent for loops. They declare the intent of the code—creating a new list—right away, rather than forcing someone to read through the logic of initializing a list, looping, and appending. They can also be slightly faster due to optimizations in the Python interpreter.

However, with great power comes responsibility. It's possible to write overly complex, nested list comprehensions that are difficult to read. A good rule of thumb is: if the logic is simple (one or two loops, maybe a simple if), a comprehension is great. If it starts to get convoluted, a standard for loop is often more maintainable.

Let's watch the final part of Corey Schafer's excellent tutorial, where he covers nested loops and also introduces dictionary and set comprehensions.

Practice Time

The best way to solidify your understanding is to write some comprehensions yourself. Try to solve the following problems. For each one, think first about how you would write it as a for loop, and then convert it to a list comprehension.

1. Given the list words = ['hello', 'world', 'in', 'python'], create a new list containing the uppercase version of each word.
   • Expected output: ['HELLO', 'WORLD', 'IN', 'PYTHON']
2. Given the same words list, create a new list containing the length of each word.
   • Expected output: [5, 5, 2, 6]
3. Create a list of numbers from 1 to 30 that are divisible by 3.
   • Expected output: [3, 6, 9, 12, 15, 18, 21, 24, 27, 30]
4. Given a list of stock data, filter for stocks with a price above $200.

   stocks = [
       {'symbol':'AAPL', 'price':120.50},
       {'symbol':'GOOG', 'price':2500.75},
       {'symbol':'META', 'price':180.25},
       {'symbol':'MSFT', 'price':280.00}
   ]
   

   • Expected output: [{'symbol': 'GOOG', 'price': 2500.75}, {'symbol': 'MSFT', 'price': 280.0}]

Conclusion

In this lesson, we've unlocked one of Python's most expressive features. You've learned how to transform traditional for loops into concise and readable list comprehensions, a skill that will make your data manipulation code much cleaner.

Key takeaways from today:

• A list comprehension provides a compact syntax for creating lists from iterables: [expression for item in iterable].
• You can add an if clause to filter items: [expression for item in iterable if condition].
• This syntax is directly analogous to mathematical set-builder notation, which makes its logic clear and powerful.
• The same concept can be extended to create dictionaries and sets.
• While powerful, comprehensions should be kept simple for readability. For complex logic, a standard for loop is often better.

You are now equipped with the fundamental tools for handling data in Python: the core data structures and a powerful way to create them. In our next lesson, we'll shift gears from procedural programming to the foundational concepts of Object-Oriented Programming (OOP) by learning how to define simple Python classes and objects. This will be essential for understanding how most Python libraries, including those for machine learning, are structured.