Allow users to define functions that shuffle vertex and edge sets

Author: weilycoder

cyaron/graph.py

@@ -46,26 +46,33 @@ def to_str(self, **kwargs):
             **kwargs(Keyword args):
                 bool shuffle = False -> whether shuffle the output or not
                 str output(Edge) = str -> the convert function which converts object Edge to str. the default way is to use str()
+                list[int] node_shuffler(list[int])
+                = lambda table: random.sample(table, k=len(table))
+                -> the random function which shuffles the vertex sequence
+                list[Edge] edge_shuffler(list[int])
+                -> a random function. the default is to shuffle the edge sequence,
+                    also, if the graph is undirected, it will swap `u` and `v` randomly.
         """
+        def _edge_shuffler_default(table):
+            edge_buf = random.sample(table, k=len(table))
+            for edge in edge_buf:
+                if not self.directed and random.randint(0, 1) == 0:
+                    (edge.start, edge.end) = (edge.end, edge.start)
+            return edge_buf
+
         shuffle = kwargs.get("shuffle", False)
         output = kwargs.get("output", str)
-        buf = []
+        node_shuffler = kwargs.get("node_shuffler", lambda table: random.sample(table, k=len(table)))
+        edge_shuffler = kwargs.get("edge_shuffler", _edge_shuffler_default)
         if shuffle:
-            new_node_id = [i for i in range(1, len(self.edges))]
-            random.shuffle(new_node_id)
-            new_node_id = [0] + new_node_id
+            new_node_id = [0] + node_shuffler(range(1, len(self.edges)))
             edge_buf = []
             for edge in self.iterate_edges():
                 edge_buf.append(
                     Edge(new_node_id[edge.start], new_node_id[edge.end], edge.weight))
-            random.shuffle(edge_buf)
-            for edge in edge_buf:
-                if not self.directed and random.randint(0, 1) == 0:
-                    (edge.start, edge.end) = (edge.end, edge.start)
-                buf.append(output(edge))
+            buf = map(output, edge_shuffler(edge_buf))
         else:
-            for edge in self.iterate_edges():
-                buf.append(output(edge))
+            buf = map(output, self.iterate_edges())
         return "\n".join(buf)
 
     def __str__(self):