Implement edge list-colouring for directed graphs.

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "vizing"
-version = "0.2.0"
+version = "0.2.0.9000"
 description = "Constraint-based list-colouring in Python."
 authors = ["Matthew Henderson <[email protected]>"]
 readme = "README.md"

src/vizing/__init__.py

@@ -27,28 +27,35 @@ def node_list_colouring_solution(G):
     G.nodes[node]['colour'] = S[node]
   return(G)
 
-def edge_list_colouring_problem(G):
+def edge_list_colouring_problem(G, directed = False):
   """
   Return a constraint problem representing an edge list-colouring instance.
 
   :param G: A graph with lists of permissible colours assigned to edges.
+  :param directed: Is the input graph directed or not?
   :return An edge list-colouring constraint problem.
   """
   P = ct.Problem()
   for edge in G.edges():
     P.addVariable(edge, G.edges[edge]['permissible'])
-  for node in G.nodes():
-    P.addConstraint(ct.AllDifferentConstraint(), [tuple(sorted(x)) for x in G.edges(node)])
+  if(directed):
+    for node in G.nodes():
+      P.addConstraint(ct.AllDifferentConstraint(), [x for x in G.out_edges(node)])
+      P.addConstraint(ct.AllDifferentConstraint(), [x for x in G.in_edges(node)])
+    else:
+      for node in G.nodes():
+        P.addConstraint(ct.AllDifferentConstraint(), [tuple(sorted(x)) for x in G.edges(node)])
   return(P)
 
-def edge_list_colouring_solution(G):
+def edge_list_colouring_solution(G, directed = False):
   """
   Return an edge-coloured graph.
 
   :param G: A graph with lists of permissible colours assigned to edges.
+  :param directed: Is the input graph directed or not?
   :return A properly edge list-coloured graph.
   """
-  P = edge_list_colouring_problem(G)
+  P = edge_list_colouring_problem(G, directed)
   S = P.getSolution()
   for edge in G.edges():
     G.edges()[edge]['colour'] = S[edge]

tests/test_edge-list-colouring.py

@@ -18,3 +18,32 @@ def test_edge_list_colouring_solution():
   assert G.edges()[(0, 1)]['colour'] == 1
   assert G.edges()[(0, 2)]['colour'] == 2
   assert G.edges()[(1, 2)]['colour'] == 4
+
+def test_directed_edge_list_colouring_problem():
+  G = nx.complete_graph(3, create_using = nx.DiGraph)
+  G.add_edge(0, 0)
+  G.add_edge(1, 1)
+  G.add_edge(2, 2)
+  permissible_dict = {(0, 0): [1], (0, 1): [2], (0, 2): [3], (1, 0): [2, 3], (1, 1): [3], (1, 2): [1], (2, 0): [2, 3], (2, 1): [1, 3], (2, 2): [2, 3]}
+  nx.set_edge_attributes(G, permissible_dict, "permissible")
+  CP = vz.edge_list_colouring_problem(G, directed = True)
+  assert CP._variables == {(0, 1): [2], (0, 2): [3], (0, 0): [1], (1, 0): [2, 3], (1, 2): [1], (1, 1): [3], (2, 0): [2, 3], (2, 1): [1, 3], (2, 2): [2, 3]}
+
+def test_directed_edge_list_colouring_solution():
+  G = nx.complete_graph(3, create_using = nx.DiGraph)
+  G.add_edge(0, 0)
+  G.add_edge(1, 1)
+  G.add_edge(2, 2)
+  permissible_dict = {(0, 0): [1], (0, 1): [2], (0, 2): [3], (1, 0): [2, 3], (1, 1): [3], (1, 2): [1], (2, 0): [2, 3], (2, 1): [1, 3], (2, 2): [2, 3]}
+  nx.set_edge_attributes(G, permissible_dict, "permissible")
+  nx.set_edge_attributes(G, None, "colour")
+  G = vz.edge_list_colouring_solution(G, directed = True)
+  assert G.edges()[(0, 0)]['colour'] == 1
+  assert G.edges()[(0, 1)]['colour'] == 2
+  assert G.edges()[(0, 2)]['colour'] == 3
+  assert G.edges()[(1, 0)]['colour'] == 2
+  assert G.edges()[(1, 1)]['colour'] == 3
+  assert G.edges()[(1, 2)]['colour'] == 1
+  assert G.edges()[(2, 0)]['colour'] == 3
+  assert G.edges()[(2, 1)]['colour'] == 1
+  assert G.edges()[(2, 2)]['colour'] == 2