Comprehensions and Lambda

By the end of this page, five-line loops will become one clear line, and you'll be able to read every pandas tutorial ever written. So much analytics code is "take this list, transform or filter it, get a new list" that Python has a compact syntax for it, the comprehension. Next to it sit lambdas, tiny anonymous functions that tell sorted, max and friends how to look at each item.

List comprehensions

prices = [300, 800, 1200, 450]

# loop version
with_gst = []
for p in prices:
    with_gst.append(p * 1.17)

# comprehension version
with_gst = [p * 1.17 for p in prices]
print(with_gst)                     # [351.0, 936.0, 1404.0, 526.5]

Read it right to left. For each p in prices, compute p * 1.17, and collect the results in a new list. The expression can be anything, from a method call to an f-string to a function you wrote.

Adding a filter

big = [p for p in prices if p > 500]
print(big)                          # [800, 1200]

names = ["  ayesha ", "BILAL", "", "usman  "]
clean = [n.strip().title() for n in names if n.strip()]
print(clean)                        # ['Ayesha', 'Bilal', 'Usman']

The if at the end decides which items get in. The second example filters (drops blanks) and transforms (tidies the case) in one line, which is the shape of most cleaning steps you'll write.

A conditional expression inside

labels = ["high" if p > 500 else "low" for p in prices]
print(labels)                       # ['low', 'high', 'high', 'low']

An if/else before the for chooses what to produce for every item. An if after the for chooses which items to keep. Students mix these up constantly, and the diagram above shows where each sits.

Dictionary and set comprehensions

products = ["Notebook", "Pen", "Marker"]
prices = [250, 40, 120]

price_of = {name: price for name, price in zip(products, prices)}
print(price_of)                     # {'Notebook': 250, 'Pen': 40, 'Marker': 120}

discounted = {name: p * 0.9 for name, p in price_of.items() if p > 100}
print(discounted)                   # {'Notebook': 225.0, 'Marker': 108.0}

cities = ["Lahore", "lahore", "Multan", "LAHORE"]
unique = {c.title() for c in cities}
print(unique)                       # {'Lahore', 'Multan'}

Nested comprehensions, sparingly

matrix = [[1, 2, 3], [4, 5, 6]]
flat = [n for row in matrix for n in row]
print(flat)                         # [1, 2, 3, 4, 5, 6]

pairs = [(b, q) for b in ["Gulberg", "DHA"] for q in ["Q1", "Q2"]]
print(pairs)

Two for clauses read in the same order as nested loops. Past two levels, or when the expression gets complicated, write a normal loop. Readability wins, and the person reviewing your pull request will agree.

Generator expressions

Swap the square brackets for parentheses and you get a generator. It produces items one at a time instead of building the whole list in memory, which is perfect when you're only feeding the items into sum, max or any.

amounts = [2500, 1800, 4100, 900]
print(sum(a for a in amounts if a > 1000))     # 8400
print(any(a > 4000 for a in amounts))          # True
print(all(a > 500 for a in amounts))           # True

With a million rows, the generator uses almost no memory. The list version would build a million-item list first and then throw it away.

Lambda, a function in one expression

gst = lambda amount: amount * 1.17
print(gst(1000))                    # 1170.0

That's the same as def gst(amount): return amount * 1.17. On its own a lambda is pointless (just use def), but it earns its place when another function needs a small rule passed in.

customers = [("Sana", 412000), ("Hamza", 98000), ("Zara", 655000)]

by_spend = sorted(customers, key=lambda c: c[1], reverse=True)
print(by_spend[0])                  # ('Zara', 655000)

top = max(customers, key=lambda c: c[1])
print(top[0])                       # Zara

words = ["banana", "fig", "apple"]
print(sorted(words, key=len))       # ['fig', 'apple', 'banana']  (key can be any function)

The key argument answers "what should I compare each item by?". A lambda keeps that answer right next to the call instead of five lines away in a def.

map and filter

amounts = ["2500", "1800", "4100"]
print(list(map(int, amounts)))                    # [2500, 1800, 4100]
print(list(filter(lambda a: a > 2000, [2500, 1800, 4100])))   # [2500, 4100]

These older tools do what comprehensions do. Most Python programmers prefer the comprehension because the intent reads more clearly, but you'll still meet map in other people's code.

Worked example, cleaning a column in one pass

raw = ["Rs 2,500", " 1800", "N/A", "Rs 4,100", ""]

def parse(text):
    try:
        return float(text.replace("Rs", "").replace(",", "").strip())
    except ValueError:
        return None

values = [parse(t) for t in raw]
print(values)                                      # [2500.0, 1800.0, None, 4100.0, None]
valid = [v for v in values if v is not None]
print(f"{len(valid)} valid, total {sum(valid):,.0f}, max {max(valid):,.0f}")

Twenty products to promote, eight lines of code: A Daraz seller in Lahore exports about 12,000 product rows with prices stored as text. One comprehension parses them, a second keeps only items priced above Rs 500 with stock on hand, and sorted(..., key=lambda r: r["margin"], reverse=True)[:20] picks the twenty highest-margin products for this week's promotion. The whole thing runs in less time than Excel takes to open the file.

Before you move on

[expr for x in items if cond] transforms and filters in one line, dict and set comprehensions use braces ({k: v for ...} and {x for ...}), and a generator expression (expr for ...) streams items into sum, max or any without building a list. lambda args: expr is a tiny function you'll mostly use as a key. And if a comprehension takes more than one reading to understand, write a loop instead.

Three to try

1. From marks = [72, 38, 91, 55, 40, 29] build a list of "PASS"/"FAIL" labels (pass mark 40) and a separate list of only the passing marks.
2. Build a dict mapping each city name to its length in characters for ["Lahore", "Karachi", "Peshawar"], then find the longest name with max and a lambda.
3. Given a list of (name, amount, days_overdue) tuples, use sorted with a lambda to order by days overdue descending then amount descending, and print the top three.

print("Hello from CSA!")

for i in range(5):
    print(i * i)