bgneal@23: """This module contains the SmallWorldGraph class for 4.4, exercise 4.
bgneal@23: 
bgneal@23: """
bgneal@23: import random
bgneal@23: 
bgneal@23: from Graph import Edge
bgneal@23: from RandomGraph import RandomGraph
bgneal@23: 
bgneal@23: 
bgneal@23: class SmallWorldGraph(RandomGraph):
bgneal@23:     """A small world graph for 4.4, exercise 4 in Think Complexity."""
bgneal@23: 
bgneal@23:     def __init__(self, vs, k, p):
bgneal@23:         """Create a small world graph. Parameters:
bgneal@23:             vs - a list of vertices
bgneal@23:             k - regular graph parameter k: the number of edges between vertices
bgneal@23:             p - probability for rewiring.
bgneal@23: 
bgneal@23:         First a regular graph is created from the list of vertices and parameter
bgneal@23:         k, which specifies how many edges each vertex should have. Then the
bgneal@23:         rewire() method is called with parameter p to rewire the graph according
bgneal@23:         to the algorithm by Watts and Strogatz.
bgneal@23: 
bgneal@23:         """
bgneal@23:         super(SmallWorldGraph, self).__init__(vs=vs)
bgneal@23:         self.add_regular_edges(k)
bgneal@23:         self.rewire(p)
bgneal@36:         self.set_edge_length(1)
bgneal@24: 
bgneal@23:     def rewire(self, p):
bgneal@23:         """Assuming the graph is initially a regular graph, rewires the graph
bgneal@23:         using the Watts and Strogatz algorithm.
bgneal@23: 
bgneal@23:         """
bgneal@23:         # We iterate over the edges in a random order:
bgneal@23:         edges = self.edges()
bgneal@23:         random.shuffle(edges)
bgneal@23:         vertices = self.vertices()
bgneal@23: 
bgneal@23:         for e in edges:
bgneal@23:             if random.random() < p:
bgneal@23:                 v, w = e    # remember the 2 vertices this edge connected
bgneal@23:                 self.remove_edge(e)     # remove from graph
bgneal@23: 
bgneal@23:                 # pick another vertex to connect to v; duplicate edges are
bgneal@23:                 # forbidden
bgneal@23:                 while True:
bgneal@23:                     x = random.choice(vertices)
bgneal@23:                     if v is not x and not self.get_edge(v, x):
bgneal@23:                         self.add_edge(Edge(v, x))
bgneal@23:                         break
bgneal@23: