Tutorial: Comprehensions, Iterators, and Iterables

Author: Florent Hivert <florent.hivert@univ-rouen.fr> and Nicolas M. Thiéry <nthiery at users.sf.net>

List comprehensions

List comprehensions are a very handy way to construct lists in Python. You can use either of the following idioms:

[ <expr> for <name> in <iterable> ]
[ <expr> for <name> in <iterable> if <condition> ]

For example, here are some lists of squares:

Sage

sage: [ i^2 for i in [1, 3, 7] ]
[1, 9, 49]
sage: [ i^2 for i in range(1,10) ]
[1, 4, 9, 16, 25, 36, 49, 64, 81]
sage: [ i^2 for i in range(1,10) if i % 2 == 1]
[1, 9, 25, 49, 81]

Python

>>> from sage.all import *
>>> [ i**Integer(2) for i in [Integer(1), Integer(3), Integer(7)] ]
[1, 9, 49]
>>> [ i**Integer(2) for i in range(Integer(1),Integer(10)) ]
[1, 4, 9, 16, 25, 36, 49, 64, 81]
>>> [ i**Integer(2) for i in range(Integer(1),Integer(10)) if i % Integer(2) == Integer(1)]
[1, 9, 25, 49, 81]

And a variant on the latter:

Sage

sage: [i^2 if i % 2 == 1 else 2 for i in range(10)]
[2, 1, 2, 9, 2, 25, 2, 49, 2, 81]

Python

>>> from sage.all import *
>>> [i**Integer(2) if i % Integer(2) == Integer(1) else Integer(2) for i in range(Integer(10))]
[2, 1, 2, 9, 2, 25, 2, 49, 2, 81]

Exercises

1. Construct the list of the squares of the prime integers between 1 and 10:

   Sage

   sage: # edit here
   

   Python

   >>> from sage.all import *
   >>> # edit here
   

2. Construct the list of the perfect squares less than 100 (hint: use srange() to get a list of Sage integers together with the method i.sqrtrem()):

   Sage

   sage: # edit here
   

   Python

   >>> from sage.all import *
   >>> # edit here
   

One can use more than one iterable in a list comprehension:

Sage

sage: [ (i,j) for i in range(1,6) for j in range(1,i) ]
[(2, 1), (3, 1), (3, 2), (4, 1), (4, 2), (4, 3), (5, 1), (5, 2), (5, 3), (5, 4)]

Python

>>> from sage.all import *
>>> [ (i,j) for i in range(Integer(1),Integer(6)) for j in range(Integer(1),i) ]
[(2, 1), (3, 1), (3, 2), (4, 1), (4, 2), (4, 3), (5, 1), (5, 2), (5, 3), (5, 4)]

Warning

Mind the order of the nested loop in the previous expression.

If instead one wants to build a list of lists, one can use nested lists as in:

Sage

sage: [ [ binomial(n, i) for i in range(n+1) ] for n in range(10) ]
[[1],
[1, 1],
[1, 2, 1],
[1, 3, 3, 1],
[1, 4, 6, 4, 1],
[1, 5, 10, 10, 5, 1],
[1, 6, 15, 20, 15, 6, 1],
[1, 7, 21, 35, 35, 21, 7, 1],
[1, 8, 28, 56, 70, 56, 28, 8, 1],
[1, 9, 36, 84, 126, 126, 84, 36, 9, 1]]

Python

>>> from sage.all import *
>>> [ [ binomial(n, i) for i in range(n+Integer(1)) ] for n in range(Integer(10)) ]
[[1],
[1, 1],
[1, 2, 1],
[1, 3, 3, 1],
[1, 4, 6, 4, 1],
[1, 5, 10, 10, 5, 1],
[1, 6, 15, 20, 15, 6, 1],
[1, 7, 21, 35, 35, 21, 7, 1],
[1, 8, 28, 56, 70, 56, 28, 8, 1],
[1, 9, 36, 84, 126, 126, 84, 36, 9, 1]]

Exercises

1. Compute the list of pairs \((i,j)\) of non negative integers such that i is at most \(5\), j is at most 8, and i and j are co-prime:

   Sage

   sage: # edit here
   

   Python

   >>> from sage.all import *
   >>> # edit here
   

2. Compute the same list for \(i < j < 10\):

   Sage

   sage: # edit here
   

   Python

   >>> from sage.all import *
   >>> # edit here
   

Iterators

Definition

To build a comprehension, Python actually uses an iterator. This is a device which runs through a bunch of objects, returning one at each call to the next method. Iterators are built using parentheses:

Sage

sage: it = (binomial(8, i) for i in range(9))
sage: next(it)
1

Python

>>> from sage.all import *
>>> it = (binomial(Integer(8), i) for i in range(Integer(9)))
>>> next(it)
1

Sage

sage: next(it)
8
sage: next(it)
28
sage: next(it)
56

Python

>>> from sage.all import *
>>> next(it)
8
>>> next(it)
28
>>> next(it)
56

You can get the list of the results that are not yet consumed:

Sage

sage: list(it)
[70, 56, 28, 8, 1]

Python

>>> from sage.all import *
>>> list(it)
[70, 56, 28, 8, 1]

Asking for more elements triggers a StopIteration exception:

Sage

sage: next(it)
Traceback (most recent call last):
...
StopIteration

Python

>>> from sage.all import *
>>> next(it)
Traceback (most recent call last):
...
StopIteration

An iterator can be used as argument for a function. The two following idioms give the same results; however, the second idiom is much more memory efficient (for large examples) as it does not expand any list in memory:

Sage

sage: sum([binomial(8, i) for i in range(9)])
256
sage: sum(binomial(8, i) for i in range(9))
256

Python

>>> from sage.all import *
>>> sum([binomial(Integer(8), i) for i in range(Integer(9))])
256
>>> sum(binomial(Integer(8), i) for i in range(Integer(9)))
256

Exercises

1. Compute the sum of \(\binom{10}{i}\) for all even \(i\):

   Sage

   sage: # edit here
   

   Python

   >>> from sage.all import *
   >>> # edit here
   

2. Compute the sum of the products of all pairs of co-prime numbers \(i, j\) for \(i<j<10\):

   Sage

   sage: # edit here
   

   Python

   >>> from sage.all import *
   >>> # edit here
   

Typical usage of iterators

Iterators are very handy with the functions all(), any(), and exists():

Sage

sage: all([True, True, True, True])
True
sage: all([True, False, True, True])
False

Python

>>> from sage.all import *
>>> all([True, True, True, True])
True
>>> all([True, False, True, True])
False

Sage

sage: any([False, False, False, False])
False
sage: any([False, False, True, False])
True

Python

>>> from sage.all import *
>>> any([False, False, False, False])
False
>>> any([False, False, True, False])
True

Let’s check that all the prime numbers larger than 2 are odd:

Sage

sage: all( is_odd(p) for p in range(1,100) if is_prime(p) and p>2 )
True

Python

>>> from sage.all import *
>>> all( is_odd(p) for p in range(Integer(1),Integer(100)) if is_prime(p) and p>Integer(2) )
True

It is well know that if 2^p-1 is prime then p is prime:

Sage

sage: def mersenne(p): return 2^p -1
sage: [ is_prime(p) for p in range(20) if is_prime(mersenne(p)) ]
[True, True, True, True, True, True, True]

Python

>>> from sage.all import *
>>> def mersenne(p): return Integer(2)**p -Integer(1)
>>> [ is_prime(p) for p in range(Integer(20)) if is_prime(mersenne(p)) ]
[True, True, True, True, True, True, True]

The converse is not true:

Sage

sage: all( is_prime(mersenne(p)) for p in range(1000) if is_prime(p) )
False

Python

>>> from sage.all import *
>>> all( is_prime(mersenne(p)) for p in range(Integer(1000)) if is_prime(p) )
False

Using a list would be much slower here:

Sage

sage: %time all( is_prime(mersenne(p)) for p in range(1000) if is_prime(p) )    # not tested
CPU times: user 0.00 s, sys: 0.00 s, total: 0.00 s
Wall time: 0.00 s
False
sage: %time all( [ is_prime(mersenne(p)) for p in range(1000) if is_prime(p)] ) # not tested
CPU times: user 0.72 s, sys: 0.00 s, total: 0.73 s
Wall time: 0.73 s
False

Python

>>> from sage.all import *
>>> %time all( is_prime(mersenne(p)) for p in range(Integer(1000)) if is_prime(p) )    # not tested
CPU times: user 0.00 s, sys: 0.00 s, total: 0.00 s
Wall time: 0.00 s
False
>>> %time all( [ is_prime(mersenne(p)) for p in range(Integer(1000)) if is_prime(p)] ) # not tested
CPU times: user 0.72 s, sys: 0.00 s, total: 0.73 s
Wall time: 0.73 s
False

You can get the counterexample using exists(). It takes two arguments: an iterator and a function which tests the property that should hold:

Sage

sage: exists( (p for p in range(1000) if is_prime(p)), lambda p: not is_prime(mersenne(p)) )
(True, 11)

Python

>>> from sage.all import *
>>> exists( (p for p in range(Integer(1000)) if is_prime(p)), lambda p: not is_prime(mersenne(p)) )
(True, 11)

An alternative way to achieve this is:

Sage

sage: counter_examples = (p for p in range(1000) if is_prime(p) and not is_prime(mersenne(p)))
sage: next(counter_examples)
11

Python

>>> from sage.all import *
>>> counter_examples = (p for p in range(Integer(1000)) if is_prime(p) and not is_prime(mersenne(p)))
>>> next(counter_examples)
11

Exercises

1. Build the list \(\{i^3 \mid -10<i<10\}\). Can you find two of those cubes \(u\) and \(v\) such that \(u + v = 218\)?

   Sage

   sage: # edit here
   

   Python

   >>> from sage.all import *
   >>> # edit here
   

itertools

At its name suggests itertools is a module which defines several handy tools for manipulating iterators:

Sage

sage: l = [3, 234, 12, 53, 23]
sage: [(i, l[i]) for i in range(len(l))]
[(0, 3), (1, 234), (2, 12), (3, 53), (4, 23)]

Python

>>> from sage.all import *
>>> l = [Integer(3), Integer(234), Integer(12), Integer(53), Integer(23)]
>>> [(i, l[i]) for i in range(len(l))]
[(0, 3), (1, 234), (2, 12), (3, 53), (4, 23)]

The same results can be obtained using enumerate():

Sage

sage: list(enumerate(l))
[(0, 3), (1, 234), (2, 12), (3, 53), (4, 23)]

Python

>>> from sage.all import *
>>> list(enumerate(l))
[(0, 3), (1, 234), (2, 12), (3, 53), (4, 23)]

Here is the analogue of list slicing:

Sage

sage: list(Permutations(3))
[[1, 2, 3], [1, 3, 2], [2, 1, 3], [2, 3, 1], [3, 1, 2], [3, 2, 1]]
sage: list(Permutations(3))[1:4]
[[1, 3, 2], [2, 1, 3], [2, 3, 1]]

sage: import itertools
sage: list(itertools.islice(Permutations(3), 1r, 4r))
[[1, 3, 2], [2, 1, 3], [2, 3, 1]]

Python

>>> from sage.all import *
>>> list(Permutations(Integer(3)))
[[1, 2, 3], [1, 3, 2], [2, 1, 3], [2, 3, 1], [3, 1, 2], [3, 2, 1]]
>>> list(Permutations(Integer(3)))[Integer(1):Integer(4)]
[[1, 3, 2], [2, 1, 3], [2, 3, 1]]

>>> import itertools
>>> list(itertools.islice(Permutations(Integer(3)), 1, 4))
[[1, 3, 2], [2, 1, 3], [2, 3, 1]]

Note that all calls to islice must have arguments of type int and not Sage integers.

The behaviour of the functions map() and filter() has changed between Python 2 and Python 3. In Python 3, they return an iterator. If you want to return a list like in Python 2 you need to explicitly wrap them in list():

Sage

sage: list(map(lambda z: z.cycle_type(), Permutations(3)))
[[1, 1, 1], [2, 1], [2, 1], [3], [3], [2, 1]]

sage: list(filter(lambda z: z.has_pattern([1,2]), Permutations(3)))
[[1, 2, 3], [1, 3, 2], [2, 1, 3], [2, 3, 1], [3, 1, 2]]

Python

>>> from sage.all import *
>>> list(map(lambda z: z.cycle_type(), Permutations(Integer(3))))
[[1, 1, 1], [2, 1], [2, 1], [3], [3], [2, 1]]

>>> list(filter(lambda z: z.has_pattern([Integer(1),Integer(2)]), Permutations(Integer(3))))
[[1, 2, 3], [1, 3, 2], [2, 1, 3], [2, 3, 1], [3, 1, 2]]

Exercises

1. Define an iterator for the \(i\)-th prime for \(5<i<10\):

   Sage

   sage: # edit here
   

   Python

   >>> from sage.all import *
   >>> # edit here
   

Defining new iterators

One can very easily write new iterators using the keyword yield. The following function does nothing interesting beyond demonstrating the use of yield:

Sage

sage: def f(n):
....:   for i in range(n):
....:       yield i
sage: [ u for u in f(5) ]
[0, 1, 2, 3, 4]

Python

>>> from sage.all import *
>>> def f(n):
...   for i in range(n):
...       yield i
>>> [ u for u in f(Integer(5)) ]
[0, 1, 2, 3, 4]

Iterators can be recursive:

Sage

sage: def words(alphabet,l):
....:    if l == 0:
....:        yield []
....:    else:
....:        for word in words(alphabet, l-1):
....:            for a in alphabet:
....:                yield word + [a]

sage: [ w for w in words(['a','b','c'], 3) ]
[['a', 'a', 'a'], ['a', 'a', 'b'], ['a', 'a', 'c'], ['a', 'b', 'a'], ['a', 'b', 'b'], ['a', 'b', 'c'], ['a', 'c', 'a'], ['a', 'c', 'b'], ['a', 'c', 'c'], ['b', 'a', 'a'], ['b', 'a', 'b'], ['b', 'a', 'c'], ['b', 'b', 'a'], ['b', 'b', 'b'], ['b', 'b', 'c'], ['b', 'c', 'a'], ['b', 'c', 'b'], ['b', 'c', 'c'], ['c', 'a', 'a'], ['c', 'a', 'b'], ['c', 'a', 'c'], ['c', 'b', 'a'], ['c', 'b', 'b'], ['c', 'b', 'c'], ['c', 'c', 'a'], ['c', 'c', 'b'], ['c', 'c', 'c']]
sage: sum(1 for w in words(['a','b','c'], 3))
27

Python

>>> from sage.all import *
>>> def words(alphabet,l):
...    if l == Integer(0):
...        yield []
...    else:
...        for word in words(alphabet, l-Integer(1)):
...            for a in alphabet:
...                yield word + [a]

>>> [ w for w in words(['a','b','c'], Integer(3)) ]
[['a', 'a', 'a'], ['a', 'a', 'b'], ['a', 'a', 'c'], ['a', 'b', 'a'], ['a', 'b', 'b'], ['a', 'b', 'c'], ['a', 'c', 'a'], ['a', 'c', 'b'], ['a', 'c', 'c'], ['b', 'a', 'a'], ['b', 'a', 'b'], ['b', 'a', 'c'], ['b', 'b', 'a'], ['b', 'b', 'b'], ['b', 'b', 'c'], ['b', 'c', 'a'], ['b', 'c', 'b'], ['b', 'c', 'c'], ['c', 'a', 'a'], ['c', 'a', 'b'], ['c', 'a', 'c'], ['c', 'b', 'a'], ['c', 'b', 'b'], ['c', 'b', 'c'], ['c', 'c', 'a'], ['c', 'c', 'b'], ['c', 'c', 'c']]
>>> sum(Integer(1) for w in words(['a','b','c'], Integer(3)))
27

Here is another recursive iterator:

Sage

sage: def dyck_words(l):
....:     if l==0:
....:         yield ''
....:     else:
....:         for k in range(l):
....:             for w1 in dyck_words(k):
....:                 for w2 in dyck_words(l-k-1):
....:                     yield '('+w1+')'+w2

sage: list(dyck_words(4))
['()()()()',
'()()(())',
'()(())()',
'()(()())',
'()((()))',
'(())()()',
'(())(())',
'(()())()',
'((()))()',
'(()()())',
'(()(()))',
'((())())',
'((()()))',
'(((())))']

sage: sum(1 for w in dyck_words(5))
42

Python

>>> from sage.all import *
>>> def dyck_words(l):
...     if l==Integer(0):
...         yield ''
...     else:
...         for k in range(l):
...             for w1 in dyck_words(k):
...                 for w2 in dyck_words(l-k-Integer(1)):
...                     yield '('+w1+')'+w2

>>> list(dyck_words(Integer(4)))
['()()()()',
'()()(())',
'()(())()',
'()(()())',
'()((()))',
'(())()()',
'(())(())',
'(()())()',
'((()))()',
'(()()())',
'(()(()))',
'((())())',
'((()()))',
'(((())))']

>>> sum(Integer(1) for w in dyck_words(Integer(5)))
42

Exercises

1. Write an iterator with two parameters \(n\), \(l\) iterating through the set of nondecreasing lists of integers smaller than \(n\) of length \(l\):

   Sage

   sage: # edit here
   

   Python

   >>> from sage.all import *
   >>> # edit here
   

Standard Iterables

Finally, many standard Python and Sage objects are iterable; that is one may iterate through their elements:

Sage

sage: sum( x^len(s) for s in Subsets(8) )
x^8 + 8*x^7 + 28*x^6 + 56*x^5 + 70*x^4 + 56*x^3 + 28*x^2 + 8*x + 1

sage: sum( x^p.length() for p in Permutations(3) )
x^3 + 2*x^2 + 2*x + 1

sage: factor(sum( x^p.length() for p in Permutations(3) ))
(x^2 + x + 1)*(x + 1)

sage: P = Permutations(5)
sage: all( p in P for p in P )
True

sage: for p in GL(2, 2): print(p); print("")
[1 0]
[0 1]

[0 1]
[1 0]

[0 1]
[1 1]

[1 1]
[0 1]

[1 1]
[1 0]

[1 0]
[1 1]


sage: for p in Partitions(3): print(p)
[3]
[2, 1]
[1, 1, 1]

Python

>>> from sage.all import *
>>> sum( x**len(s) for s in Subsets(Integer(8)) )
x^8 + 8*x^7 + 28*x^6 + 56*x^5 + 70*x^4 + 56*x^3 + 28*x^2 + 8*x + 1

>>> sum( x**p.length() for p in Permutations(Integer(3)) )
x^3 + 2*x^2 + 2*x + 1

>>> factor(sum( x**p.length() for p in Permutations(Integer(3)) ))
(x^2 + x + 1)*(x + 1)

>>> P = Permutations(Integer(5))
>>> all( p in P for p in P )
True

>>> for p in GL(Integer(2), Integer(2)): print(p); print("")
[1 0]
[0 1]
<BLANKLINE>
[0 1]
[1 0]
<BLANKLINE>
[0 1]
[1 1]
<BLANKLINE>
[1 1]
[0 1]
<BLANKLINE>
[1 1]
[1 0]
<BLANKLINE>
[1 0]
[1 1]
<BLANKLINE>

>>> for p in Partitions(Integer(3)): print(p)
[3]
[2, 1]
[1, 1, 1]

Beware of infinite loops:

Sage

sage: for p in Partitions(): print(p)          # not tested

Python

>>> from sage.all import *
>>> for p in Partitions(): print(p)          # not tested

Sage

sage: for p in Primes(): print(p)              # not tested

Python

>>> from sage.all import *
>>> for p in Primes(): print(p)              # not tested

Infinite loops can nevertheless be very useful:

Sage

sage: exists( Primes(), lambda p: not is_prime(mersenne(p)) )
(True, 11)


sage: counter_examples = (p for p in Primes() if not is_prime(mersenne(p)))
sage: next(counter_examples)
11

Python

>>> from sage.all import *
>>> exists( Primes(), lambda p: not is_prime(mersenne(p)) )
(True, 11)


>>> counter_examples = (p for p in Primes() if not is_prime(mersenne(p)))
>>> next(counter_examples)
11