generators_soln.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)

(0, 1)
(0, 2)
(0, 3)
(1, 2)
(1, 3)
(2, 3)

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

(0, 1)
(0, 3)
(2, 3)

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

# Solution

def random_pairs(nodes, p):
    for edge in all_pairs(nodes):
        if flip(p):
            yield edge

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

(0, 2)
(1, 3)

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

from random import randrange

def make_random_tree(n):
    G = nx.Graph()
    if n == 0:
        return G

    G.add_node(0)
    for v in range(1, n):
        u = randrange(v)
        G.add_edge(u, v)

    return G

tree = make_random_tree(10)
tree.nodes()

NodeView((0, 1, 2, 3, 4, 5, 6, 7, 8, 9))

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

def all_neighbors(G):
    for u in G.nodes():
        for v in G.neighbors(u):
            yield u, v

# Solution

def all_triangles(G):
    for u, v in all_neighbors(G):
        ws = set(G.neighbors(u)) & set(G.neighbors(v))
        for w in ws:
            yield u, v, w

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

NodeView((0, 1, 2))

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

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

(0, 1, 2)
(0, 2, 1)
(1, 0, 2)
(1, 2, 0)
(2, 0, 1)
(2, 1, 0)

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