generators.ipynb

Kernel: Python 3

Generator functions

Code examples from Think Complexity, 2nd edition.

Copyright 2019 Allen Downey, MIT License

%matplotlib inline

import networkx as nx
import numpy as np

# TODO: remove this when NetworkX is fixed
from warnings import simplefilter
import matplotlib.cbook
simplefilter("ignore", matplotlib.cbook.mplDeprecation)

Exercise: In the book I wrote a version of random_pairs that violates Ned’s recommendation to “abstract your iteration”:

def flip(p):
    return np.random.random() < p

def all_pairs(nodes):
    for i, u in enumerate(nodes):
        for j, v in enumerate(nodes):
            if i < j:
                yield u, v

def random_pairs(nodes, p):
    for i, u in enumerate(nodes):
        for j, v in enumerate(nodes):
            if i < j and flip(p):
                yield u, v

nodes = range(4)

for pair in all_pairs(nodes):
    print(pair)

for pair in random_pairs(nodes, 0.5):
    print(pair)

Write a better version of this function that uses all_pairs rather than copying and modifying it.

# Solution goes here

for pair in random_pairs(nodes, 0.5):
    print(pair)

Exercise: Write a function called random_tree that takes a number of nodes, n, as a parameter and builds an undirected graph by starting with a single node, adding one node at a time, and connecting each new node to one existing node. You can use any of the functions in Python’s random module.

# Solution goes here

tree = make_random_tree(10)
tree.nodes()

nx.draw(tree, 
        node_color='C0', 
        node_size=1000, 
        with_labels=True)

Bonus: Read the various equivalent definitions of tree and then write a function called is_tree that takes a graph and returns True if the graph is a tree.

Exercise: Write a function called all_triangles that takes an undirected graph as a parameter and returns all triangles, where a triangle is a collection of three nodes that are connected to each other (regardless of whether they are also connected to other nodes). Your solution can be an ordinary function that returns a list of tuples, or a generator function that yields tuples. It does not have to be particularly efficient. It’s OK if your solution finds the same triangle more than once, but as a bonus challenge, write a solution that avoids it.

# Solution goes here

# Solution goes here

def make_complete_graph(n):
    G = nx.Graph()
    nodes = range(n)
    G.add_nodes_from(nodes)
    G.add_edges_from(all_pairs(nodes))
    return G

complete = make_complete_graph(3)
complete.nodes()

nx.draw_circular(complete, 
                 node_color='C1', 
                 node_size=1000, 
                 with_labels=True)

for tri in all_triangles(complete):
    print(tri)

for tri in all_triangles(tree):
    print(tri)