Python's More Advanced Features: Part I - Iterators, Generators, and itertools

collatz.py

def collatz(n):
    m = n
    while True:
        yield m
        if m == 1:
            raise StopIteration
        if m % 2 == 0:
            m /= 2
        else:
            m = 3 * m + 1


# If the Collatz conjecture is true then this program will
# always terminate.
c = collatz(7)
for m in c:
    print m


# Output:
7
22
11
34
17
52
26
13
40
20
10
5
16
8
4
2
1

count_generator.py

def count_generator():
    n = 0
    while True:
        yield n
        n += 1

count_iterator.py

class count_iterator(object):
    n = 0

    def __iter__(self):
        return self

    def next(self):
        y = self.n
        self.n += 1
        return y

derangements.py

def count_fixed_points(p):
    """Return the number of fixed points of p as a permutation."""
    return sum(1 for x in p if p[x]==x)

# Note: This is NOT an efficient way of doing this. 
# Just to demonstrate the use of itertools and generators
def count_partial_derangements(n,k):
    """Returns the number of permutations of n with k fixed points."""
    return sum(1 for p in itertools.permutations(xrange(n)) if count_fixed_points(p) == k) 

# Usage:
>>> [count_partial_derangements(6,i) for i in xrange(7)]
[265, 264, 135, 40, 15, 0, 1]

generator_expression_1.py

>>> g = (x**2 for x in itertools.count(1))
>>> [g.next() for __ in xrange(10)]
[1, 4, 9, 16, 25, 36, 49, 64, 81, 100]

generator_expression_2.py

>>> g = (random.random()<0.4 for __ in itertools.count())
>>> [g.next() for __ in xrange(10)]
[False, False, False, True, True, False, True, False, False, True]

hofstadter_generator.py

def hofstadter_generator(n):
    a = s[:]
    while True:
        try:
            q = a[-a[-1]] + a[-a[-2]]
            a.append(q)
            yield q
        except IndexError:
            raise StopIteration()

hofstadter_iterator.py

class hofstadter_iterator(object):
    def __init__(self, s):
        self.s = s[:]

    def current_state(self):
        return self.s

    def next(self):
        try:
            q = self.s[-self.s[-1]] + self.s[-self.s[-2]]
            self.s.append(q)
            return q
        except IndexError:
            raise StopIteration()

    def fast_forward(self, n):
        for __ in range(n):
            self.next()

    def __iter__(self):
        return self

inversion_number.py

import itertools
import math


def inversion_number(A):
    """Return the number of inversions in list A."""
    return sum(1 for x, y in itertools.combinations(xrange(len(A)), 2) if A[x] > A[y])


def total_inversions(n):
    """Return total number of inversions in permutations of n."""
    return sum(inversion_number(A) for A in itertools.permutations(xrange(n)))

# Usage:
>>> [total_inversions(n) for n in xrange(10)]
[0, 0, 1, 9, 72, 600, 5400, 52920, 564480, 6531840]

>>> [math.factorial(n)*n*(n-1)/4 for n in xrange(10)]
[0, 0, 1, 9, 72, 600, 5400, 52920, 564480, 6531840]

tent_map.py

def tent_map(mu, x0):
    x = x0
    while True:
        yield x
        x = mu * min(x, 1.0 - x)


t = tent_map(2.0, 0.1)
for __ in xrange(30):
    print t.next()


# Output: 
0.1
0.2
0.4
0.8
0.4
0.8
0.4
0.8
0.4
0.8
0.4
0.8
0.4
0.8
0.4
0.8
0.4
0.799999999999
0.400000000001
0.800000000003
0.399999999994
0.799999999988
0.400000000023
0.800000000047
0.399999999907
0.799999999814
0.400000000373
0.800000000745
0.39999999851
0.79999999702