Implement total list colouring

pyproject.toml

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

src/vizing/__init__.py

@@ -1,4 +1,4 @@
-import constraint
+import constraint as ct
 
 def node_list_colouring_problem(G):
   """
@@ -7,11 +7,11 @@ def node_list_colouring_problem(G):
   :param G: A graph with lists of permissible colours assigned to nodes.
   :return A node list-colouring constraint problem.
   """
-  P = constraint.Problem()
+  P = ct.Problem()
   for node in G.nodes():
     P.addVariable(node, G.nodes[node]['permissible'])
   for edge in G.edges():
-    P.addConstraint(constraint.AllDifferentConstraint(), edge)
+    P.addConstraint(ct.AllDifferentConstraint(), edge)
   return(P)
 
 def node_list_colouring_solution(G):
@@ -34,11 +34,11 @@ def edge_list_colouring_problem(G):
   :param G: A graph with lists of permissible colours assigned to edges.
   :return An edge list-colouring constraint problem.
   """
-  P = constraint.Problem()
+  P = ct.Problem()
   for edge in G.edges():
     P.addVariable(edge, G.edges[edge]['permissible'])
   for node in G.nodes():
-    P.addConstraint(constraint.AllDifferentConstraint(), [tuple(sorted(x)) for x in G.edges(node)])
+    P.addConstraint(ct.AllDifferentConstraint(), [tuple(sorted(x)) for x in G.edges(node)])
   return(P)
 
 def edge_list_colouring_solution(G):
@@ -53,3 +53,37 @@ def edge_list_colouring_solution(G):
   for edge in G.edges():
     G.edges()[edge]['colour'] = S[edge]
   return(G)
+
+def total_list_colouring_problem(G):
+  """
+  Return a constraint problem representing a total list-colouring instance.
+
+  :param G: A graph with lists of permissible colours assigned to nodes and edges.
+  :return A total list-colouring constraint problem.
+  """
+  P = ct.Problem()
+  for node in G.nodes():
+    P.addVariable(node, G.nodes[node]['permissible'])
+  for edge in G.edges():
+    P.addVariable(edge, G.edges[edge]['permissible'])
+  for edge in G.edges():
+    P.addConstraint(ct.AllDifferentConstraint(), [edge[0], edge[1], edge])
+  for node in G.nodes():
+    P.addConstraint(ct.AllDifferentConstraint(), [node] + [tuple(sorted(x)) for x in G.edges(node)])
+  return(P)
+
+def total_list_colouring_solution(G):
+  """
+  Return a total list-coloured graph.
+
+  :param G: A graph with lists of permissible colours assigned to nodes and edges.
+  :return A properly total list-coloured graph.
+  """
+  P = total_list_colouring_problem(G)
+  S = P.getSolution()
+  for node in G.nodes:
+    G.nodes[node]['colour'] = S[node]
+  for edge in G.edges():
+    G.edges()[edge]['colour'] = S[edge]
+  return(G)
+

tests/test_total-list-colouring.py

@@ -0,0 +1,23 @@
+import listcolouring as lc
+import networkx as nx
+import vizing as vz
+
+def test_total_list_colouring_problem():
+  G = nx.complete_graph(3)
+  permissible_dict_edge = {(0, 1): [0, 1], (0, 2): [1, 2], (1, 2): [2, 4]}
+  nx.set_edge_attributes(G, permissible_dict_edge, "permissible")
+  nx.set_edge_attributes(G, None, "colour")
+  permissible_dict_node = {0: [0, 1], 1: [1, 2], 2: [2, 4]}
+  nx.set_node_attributes(G, permissible_dict_node, "permissible")
+  nx.set_node_attributes(G, None, "colour")
+  P = vz.total_list_colouring_problem(G)
+  assert P._variables == {0: [0, 1], 1: [1, 2], 2: [2, 4], (0, 1): [0, 1], (0, 2): [1, 2], (1, 2): [2, 4]}
+
+def test_total_list_colouring_solution():
+  G = nx.petersen_graph()
+  n_colours = 6
+  G = lc.list_init_node(G, range(n_colours - 1), 3, 0)
+  G = lc.list_init(G, range(n_colours - 1), 3, 0)
+  G = vz.total_list_colouring_solution(G)
+  assert G.nodes()[0]['colour'] == 0
+