Comprehensions

lecture
python
comprehensions
for-loops

See csse1001 for course logistics — this note covers Lecture 5B’s technical content. See python-comprehensions for the full reference on comprehension syntax.

Today’s outline

  • List comprehensions
  • Dictionary comprehensions
  • Nested comprehensions
  • Filtering
  • all and any

List comprehension

Mathematics builds sets and sums this way:

\[ \{k^2 : k \in \{0, 1, 2, 3\}\} = \{0, 1, 4, 9\} \qquad \sum_{k \in \{0,1,2,3\}} k^2 = 0^2+1^2+2^2+3^2 \]

Python mirrors this directly:

>>> [k**2 for k in range(4)]
[0, 1, 4, 9]
>>> sum(k**2 for k in range(4))
14

Dictionary comprehension

>>> {x: f"value is {x}" for x in range(3)}
{0: 'value is 0', 1: 'value is 1', 2: 'value is 2'}

>>> {x: 0 for x in "ABC"}   # useful for assigning default values
{'A': 0, 'B': 0, 'C': 0}

>>> cs = "hello world"
>>> {x: cs.count(x) for x in cs}
{'h': 1, 'e': 1, 'l': 3, 'o': 2, ' ': 1, 'w': 1, 'r': 1, 'd': 1}

Nested comprehensions

Multiple for clauses can appear in one comprehension — order matters:

>>> [(a, x) for a in "ABC" for x in range(2)]
[('A', 0), ('A', 1), ('B', 0), ('B', 1), ('C', 0), ('C', 1)]

>>> [(a, x) for x in range(2) for a in "ABC"]
[('A', 0), ('B', 0), ('C', 0), ('A', 1), ('B', 1), ('C', 1)]

Implementing [(a, x) for a in "ABC" for x in range(2)] with a for-loop:

>>> xs = []
>>> for a in "ABC":
...     for x in range(2):
...         xs.append((a, x))

Exercise: disassemble a string

def disassemble(cs: str) -> list[str]:
    """ Returns a list of characters comprising <cs>, maintaining order.
    >>> disassemble("")
    []
    >>> disassemble("abcdef")
    ['a', 'b', 'c', 'd', 'e', 'f']
    """

With a for-loop:

>>> def disassemble(xs: str) -> list[str]:
...     ans = []
...     for x in xs:
...         ans.append(x)
...     return ans

With a list comprehension:

The companion source file defines disassemble with the comprehension return first, followed by an equivalent for-loop version left in the function body below it. That second block is unreachable dead code (the return above it always exits first) and is omitted here.

return [c for c in cs]

Exercise: more capitals than lower-case

Check if a string has (strictly) more capitals than lower-case letters, ignoring all other characters:

def more_capitals(cs: str) -> bool:
    """
    >>> more_capitals("aAbBcCdD")
    False
    >>> more_capitals("CSSE 1001")
    True
    >>> more_capitals("")
    False
    """

A useful “hack”: the integer value of True and False is 1 and 0 respectively, so sum(...) over a sequence of booleans counts how many are True:

>>> True + True + False
2

The lecture slide’s solution has a typo — it compares against 'a' <= c <= 'a' (matching only the literal character 'a') instead of 'a' <= c <= 'z' (matching any lower-case letter). The companion source file has it correct; the exercise below uses the corrected bound.

return (sum('A' <= c <= 'Z' for c in cs)
        > sum('a' <= c <= 'z' for c in cs))

Exercise: all

Check if every element of a list is truthy:

def all(xs: list) -> bool:
    """
    >>> all([1, 2, 3])
    True
    >>> all(["a", "", "c"])
    False
    >>> all([])
    True
    """
return len(xs) == sum(bool(x) for x in xs)

Exercise: any

Check if any element of a list is truthy:

def any(xs: list) -> bool:
    """
    >>> any([0, 0, 0])
    False
    >>> any(["a", "", "c"])
    True
    >>> any([])
    False
    """
return sum(bool(x) for x in xs) > 0

all and any are both included in base Python already (no import needed):

>>> xs = "HELLO WORLD"
>>> all('A' <= x <= 'Z' for x in xs)   # is every letter a capital?
False    # (space is not a capital)
>>> any('A' <= x <= 'Z' for x in xs)   # is any letter a capital?
True

Filtering

Comprehension elements may be filtered with a trailing if. E.g. the numbers less than 100 divisible by both 3 and 7:

>>> [x for x in range(101) if not (x % 3) and not (x % 7)]
[0, 21, 42, 63, 84]

The general pattern

[x for x in xs if P(x)]
[x for x in xs for y in ys if P(x, y)]
[x for x in xs for y in ys for z in zs if P(x, y, z)]
...

where P is some predicate.

Summary

Comprehensions are a rapid way of forming lists and dictionaries that would otherwise require an accumulator for-loop.

Next: week-five-exercises — practice exercises from Lecture 5C.