.. _tutorial-comprehensions:

==================================================
Tutorial: Comprehensions, Iterators, and Iterables
==================================================

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

.. linkall

List comprehensions
===================

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

.. CODE-BLOCK:: python

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

For example, here are some lists of squares::

    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]

And a variant on the latter::

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

.. TOPIC:: Exercises

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

            sage: # edit here

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

            sage: # edit here

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

    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)]

.. 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: [ [ 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]]

.. TOPIC:: Exercises

    #.  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: # edit here

    #.  Compute the same list for `i < j < 10`::

            sage: # 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: it = (binomial(8, i) for i in range(9))
    sage: next(it)
    1

::

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

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

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

Asking for more elements triggers a ``StopIteration`` exception::

    sage: 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: sum([binomial(8, i) for i in range(9)])
    256
    sage: sum(binomial(8, i) for i in xrange(9))  # py2
    256
    sage: sum(binomial(8, i) for i in range(9))  # py3
    256

.. TOPIC:: Exercises

    #.  Compute the sum of `\binom{10}{i}` for all even `i`::

            sage: # edit here

    #.  Compute the sum of the products of all pairs of co-prime numbers `i, j` for `i<j<10`::

            sage: # edit here


Typical usage of iterators
--------------------------

Iterators are very handy with the functions :func:`all`, :func:`any`,
and :func:`exists`::

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

::

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

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

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

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

    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]

The converse is not true::

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

Using a list would be much slower here::

    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

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

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

An alternative way to achieve this is::

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

.. TOPIC:: Exercises

    #.  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: # edit here

itertools
---------

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

    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)]

The same results can be obtained using :func:`enumerate`::

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

Here is the analogue of list slicing::

    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]]

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

The behaviour of the functions :func:`map` and :func:`filter` has
changed between Python 2 and Python 3. In Python 3, they return an
iterator. If you want to use this new behaviour in Python 2, and keep
your code compatible with Python3, you can use the compatibility
library ``six`` as follows::

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

    sage: from six.moves import filter
    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]]

.. TOPIC:: Exercises

    #.  Define an iterator for the `i`-th prime for `5<i<10`::

           sage: # 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: def f(n):
    ....:   for i in range(n):
    ....:       yield i
    sage: [ u for u in f(5) ]
    [0, 1, 2, 3, 4]

Iterators can be recursive::

    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

Here is another recursive iterator::

    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

.. TOPIC:: Exercises

    #.  Write an iterator with two parameters `n`, `l` iterating
        through the set of nondecreasing lists of integers smaller
        than `n` of length `l`::

           sage: # edit here


Standard Iterables
==================

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

    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]
    <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>

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


.. skip

Beware of infinite loops::

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

.. skip

::

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

Infinite loops can nevertheless be very useful::

    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