Skip to main content

Python Comprehensions

Python comprehensions are elegant, concise ways to create or transform collections like lists, dictionaries, and sets in a single line of code. They are one of Python's most distinctive and powerful features that you'll use frequently when working with Pandas and other data processing tasks.

What are Comprehensions?​

Comprehensions provide a compact syntax to generate collections from other collections or iterables. They combine the operations of:

  • Looping through sequences
  • Conditional filtering
  • Creating new elements

Instead of writing multiple lines with loops and conditional statements, you can express these operations in a single, readable line.

Types of Comprehensions​

Python offers three main types of comprehensions:

  1. List comprehensions - Create lists
  2. Dictionary comprehensions - Create dictionaries
  3. Set comprehensions - Create sets

Let's explore each type in detail.

List Comprehensions​

List comprehensions create new lists by applying an expression to each item in an iterable.

Basic Syntax​

python
[expression for item in iterable]

Simple Example​

python
# Traditional way
squares = []
for i in range(1, 6):
    squares.append(i**2)
print(squares)

# Using list comprehension
squares = [i**2 for i in range(1, 6)]
print(squares)

Output:

[1, 4, 9, 16, 25]
[1, 4, 9, 16, 25]

List Comprehension with Conditional Logic​

You can add conditions to filter elements:

python
# Get only even squares
even_squares = [i**2 for i in range(1, 11) if i % 2 == 0]
print(even_squares)

Output:

[4, 16, 36, 64, 100]

You can also use if-else conditions within comprehensions:

python
# Mark numbers as even or odd
number_types = ["even" if x % 2 == 0 else "odd" for x in range(5)]
print(number_types)

Output:

['even', 'odd', 'even', 'odd', 'even']

Dictionary Comprehensions​

Dictionary comprehensions create dictionaries using a similar syntax, but with key-value pairs.

Basic Syntax​

python
{key_expression: value_expression for item in iterable}

Simple Example​

python
# Create a dictionary of number:square pairs
number_squares = {x: x**2 for x in range(1, 6)}
print(number_squares)

Output:

{1: 1, 2: 4, 3: 9, 4: 16, 5: 25}

Dictionary Comprehension with Conditional Logic​

python
# Create a dictionary with only even number squares
even_squares_dict = {x: x**2 for x in range(1, 11) if x % 2 == 0}
print(even_squares_dict)

Output:

{2: 4, 4: 16, 6: 36, 8: 64, 10: 100}

Set Comprehensions​

Set comprehensions create sets (collections with unique elements) using a similar syntax.

Basic Syntax​

python
{expression for item in iterable}

Simple Example​

python
# Create a set of squares
square_set = {x**2 for x in range(1, 6)}
print(square_set)

Output:

{1, 4, 9, 16, 25}

Note that the output order might vary since sets are unordered collections.

With Conditional Logic​

python
# Create a set of even squares
even_square_set = {x**2 for x in range(1, 11) if x % 2 == 0}
print(even_square_set)

Output:

{4, 16, 36, 64, 100}

Nested Comprehensions​

You can create more complex structures with nested comprehensions:

python
# Create a matrix (list of lists) using nested comprehensions
matrix = [[i*j for j in range(1, 4)] for i in range(1, 4)]
print(matrix)

Output:

[[1, 2, 3], [2, 4, 6], [3, 6, 9]]

Practical Applications for Data Work​

Comprehensions are especially useful when working with data. Here are some practical examples:

Data Cleaning​

python
# Remove None values from a list
data = [1, None, 3, None, 5, 6]
cleaned_data = [x for x in data if x is not None]
print(cleaned_data)

Output:

[1, 3, 5, 6]

Data Transformation​

python
# Convert temperatures from Celsius to Fahrenheit
celsius = [0, 10, 20, 30, 40]
fahrenheit = [(9/5) * c + 32 for c in celsius]
print(fahrenheit)

Output:

[32.0, 50.0, 68.0, 86.0, 104.0]

Creating a Lookup Dictionary​

python
# Create a dictionary from two lists
fruits = ['apple', 'banana', 'orange', 'mango']
colors = ['red', 'yellow', 'orange', 'yellow']
fruit_colors = {fruit: color for fruit, color in zip(fruits, colors)}
print(fruit_colors)

Output:

{'apple': 'red', 'banana': 'yellow', 'orange': 'orange', 'mango': 'yellow'}

Flattening a 2D List​

python
# Convert a 2D list into a flat list
matrix = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
flattened = [num for row in matrix for num in row]
print(flattened)

Output:

[1, 2, 3, 4, 5, 6, 7, 8, 9]

Performance Considerations​

Comprehensions are not only more concise but also typically faster than equivalent code using for loops. They're optimized at the bytecode level in Python.

However, be careful with very large datasets or complex expressions:

python
# This is fine for small ranges
[x**2 for x in range(1000)]

# For very large ranges, consider using a generator expression instead
# (which doesn't store all results in memory at once)
sum(x**2 for x in range(1000000))

When to Use Comprehensions​

Comprehensions are best used when:

  • You need to create a new collection based on an existing one
  • The transformation logic is simple and fits on one line
  • You want to combine filtering and transformation

Consider using regular loops when:

  • The logic is complex or hard to read in a single line
  • You need to handle exceptions
  • You're performing operations without creating a new collection

Summary​

Python comprehensions provide a powerful, concise way to create and transform collections. They combine loops, conditional logic, and collection creation into a single, readable expression. The three main types are:

  1. List comprehensions - Create new lists
  2. Dictionary comprehensions - Create key-value pair dictionaries
  3. Set comprehensions - Create sets with unique elements

Mastering comprehensions will significantly improve your Python code's readability and efficiency, especially when working with data in Pandas.

Exercises​

  1. Convert this loop to a list comprehension:

    python
    words = ["hello", "world", "python", "comprehension"]
    uppercase_words = []
    for word in words:
        uppercase_words.append(word.upper())

  2. Create a dictionary comprehension that takes a list of strings and maps each string to its length.

  3. Use a list comprehension to extract all vowels from a sentence.

  4. Create a set comprehension that collects all unique even digits from a number.

  5. Write a nested list comprehension to transpose a matrix (swap rows and columns).

Additional Resources​

Happy coding with comprehensions!


If you spot any mistakes on this website, please let me know at [email protected]. I’d greatly appreciate your feedback! :)