Comprehension

From CS 61A Wiki
(Redirected from List comprehension)
Jump to: navigation, search

A comprehension is a construct that allows sequences to be built from other sequences.

List comprehension

A list comprehension is a concise way to create a list. The syntax is: [expression for elem in sequence if condition]. The conditional is optional. This is equivalent to lst in the following code:

lst = []
for elem in sequence:
    if condition:
        lst.append(expression)

Examples:

>>> def map(f, seq):
...     return [f(elem) for elem in seq]
...
>>> map(lambda x: x * x, [1, 2, 3])
[1, 4, 9]
>>> def filter(f, seq):
...     return [elem for elem in seq if f(elem)]
...
>>> filter(lambda x: x % 2 == 0, [1, 2, 3, 4])
[2, 4]

A list comprehension can have multiple for statements. The syntax is: [expression for x in sequence1 for y in sequence2]. The loops unroll in the same order as in the comprehension:

lst = []
for x in sequence1:
    for y in sequence2:
        lst.append(expression)

Example:

>>> problems = [1, 2, 3, 4, 5, 6]
>>> parts = 'abcdef'
>>> [str(a) + b for a in problems for b in parts] # Sahai's CS70 problem sets (context: http://inst.eecs.berkeley.edu/~cs70/sp14/)
['1a', '1b', '1c', '1d', '1e', '1f', '2a', '2b', '2c', '2d', '2e', '2f', '3a', '3b', '3c', '3d', '3e', '3f', '4a', '4b', '4c', '4d', '4e', '4f', '5a', '5b', '5c', '5d', '5e', '5f', '6a', '6b', '6c', '6d', '6e', '6f']

Dictionary comprehension

A dictionary comprehension is a concise way to create a dictionary. The syntax is {expr1 : expr2 for elem in sequence if condition}. The conditional is optional. This is equivalent to d in the following code:

d = dict([(expr1, expr2) for elem in sequence if condition]) # using dict constructor; see article "dictionary"

or

d = {}
for elem in sequence:
    if condition:
        d[expr1] = expr2

Example:

>>> d = {'anDreW' : 'HuAng', 'ROHiN' : 'ShaH'}
>>> {k.title() : v.title() for k, v in d.items()} # fix capitalization
{'Andrew': 'Huang', 'Rohin': 'Shah'}

See also

Sources