Skip to main content

List Comprehensions

In Python, there is a short and powerful way to create new lists from existing data — it is called a List Comprehension. It does the same thing as a for loop but in just one line.


1. The Problem: Too Many Lines for Simple Tasks

Imagine you want to create a list of squares from 1 to 5.

The Long Way (Using a Loop):

squares = []
for num in range(1, 6):
squares.append(num ** 2)

print(squares)
# Output: [1, 4, 9, 16, 25]

That works, but it takes 4 lines for a simple task.

The Short Way (List Comprehension):

squares = [num ** 2 for num in range(1, 6)]

print(squares)
# Output: [1, 4, 9, 16, 25]

Same result, just 1 line! This is a list comprehension.


2. The Basic Formula

new_list = [expression for item in iterable]
  • expression → What you want to do with each item (e.g., num ** 2)
  • item → The variable name for each element (e.g., num)
  • iterable → The data you are looping through (e.g., range(1, 6) or a list)

More Examples:

# Double every number
numbers = [1, 2, 3, 4, 5]
doubled = [n * 2 for n in numbers]
print(doubled) # Output: [2, 4, 6, 8, 10]

# Convert names to uppercase
names = ["sai", "ravi", "priya"]
upper_names = [name.upper() for name in names]
print(upper_names) # Output: ['SAI', 'RAVI', 'PRIYA']

# Get the length of each word
words = ["python", "ai", "ml"]
lengths = [len(word) for word in words]
print(lengths) # Output: [6, 2, 2]

3. Adding a Condition (Filtering)

You can add an if condition to pick only certain items:

new_list = [expression for item in iterable if condition]

Examples:

# Only even numbers
numbers = [1, 2, 3, 4, 5, 6, 7, 8]
evens = [n for n in numbers if n % 2 == 0]
print(evens) # Output: [2, 4, 6, 8]

# Only names with more than 3 letters
names = ["sai", "ravi", "priya", "raj"]
long_names = [name for name in names if len(name) > 3]
print(long_names) # Output: ['ravi', 'priya']

# Squares of only odd numbers
odd_squares = [n ** 2 for n in range(1, 11) if n % 2 != 0]
print(odd_squares) # Output: [1, 9, 25, 49, 81]

4. If-Else Inside a Comprehension

When you want to do different things based on a condition:

# Mark "pass" or "fail" based on marks
marks = [85, 40, 72, 30, 91]
results = ["pass" if m >= 50 else "fail" for m in marks]
print(results)
# Output: ['pass', 'fail', 'pass', 'fail', 'pass']

Note: When using if-else together, the condition goes before the for. When using only if (filtering), it goes after the for.


5. Nested List Comprehensions

You can also loop through 2D data (a list inside a list):

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

Tip: Read nested comprehensions from left to right — the outer loop comes first.


6. Why This Matters for AI & ML

In AI and ML code, you will see list comprehensions everywhere:

# Clean a dataset: Remove empty strings
raw_data = ["hello", "", "world", "", "python"]
clean_data = [item for item in raw_data if item != ""]
print(clean_data) # Output: ['hello', 'world', 'python']

# Extract specific features from data
students = [
{"name": "Sai", "score": 85},
{"name": "Ravi", "score": 72},
{"name": "Priya", "score": 91},
]
scores = [student["score"] for student in students]
print(scores) # Output: [85, 72, 91]

Quick Summary

  • List Comprehension: A concise, one-line syntax for creating transformed lists from loops.
  • Basic Formula: [expression for item in iterable].
  • Filtered Formula: [expression for item in iterable if condition].
  • Conditional Transformation: [value_if_true if condition else value_if_false for item in iterable].
  • AI & Data Science: Extensively used for data preprocessing, tokenization, and feature extraction.

What's Next?

Let's explore how to create dictionaries and sets dynamically in a single line using Dictionary & Set Comprehensions in the next lesson!