From 8043b0aed362f96a5981d5a0b2f953669627aeca Mon Sep 17 00:00:00 2001
From: Matthew Henderson <matthew.james.henderson@gmail.com>
Date: Tue, 19 Nov 2024 10:20:54 +0000
Subject: [PATCH] Add implementation of greedy list total colouring. (#39)

---
 listcolouring/__init__.py                 | 15 +++++++++++----
 tests/test_greedy-list-total-colouring.py | 16 ++++++++++++++++
 2 files changed, 27 insertions(+), 4 deletions(-)
 create mode 100644 tests/test_greedy-list-total-colouring.py

diff --git a/listcolouring/__init__.py b/listcolouring/__init__.py
index 56e91b6..d371549 100644
--- a/listcolouring/__init__.py
+++ b/listcolouring/__init__.py
@@ -77,11 +77,18 @@ def greedy_list_edge_colouring(G):
         G[u][v]["colour"] = first_permissible_or_none(G, u, v) # random.choice(colours)
     return(G)
 
-def greedy_list_node_colouring(G):
-    """Assign the first permissible colour to every node (or None if all permissible
-    colours already used on incident edges)."""
+def greedy_list_node_colouring(G, edges = False):
+    """Assign the first permissible colour to every node or None if all permissible
+    colours already used on neighbouring nodes (and incident edges if edges = True)."""
     for u in G.nodes:
-        G.nodes[u]["colour"] = first_permissible_or_none_node(G, u)
+        G.nodes[u]["colour"] = first_permissible_or_none_node(G, u, edges)
+    return(G)
+
+def greedy_list_total_colouring(G):
+    """Use a greedy strategy to colour the edges of G and then use a greedy node colouring
+    strategy to colour the nodes of G."""
+    G = greedy_list_edge_colouring(G)
+    G = greedy_list_node_colouring(G, edges = True)
     return(G)
 
 def print_list_edge_colouring(G):
diff --git a/tests/test_greedy-list-total-colouring.py b/tests/test_greedy-list-total-colouring.py
new file mode 100644
index 0000000..7928923
--- /dev/null
+++ b/tests/test_greedy-list-total-colouring.py
@@ -0,0 +1,16 @@
+import networkx as nx
+
+import listcolouring
+from listcolouring import greedy_list_total_colouring
+
+def test_greedy_list_total_colouring():
+    G = nx.complete_graph(3)
+    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")
+    nx.set_edge_attributes(G, [1, 2, 3], "permissible")
+    nx.set_edge_attributes(G, 2, "colour")
+    G = greedy_list_total_colouring(G)
+    assert G.nodes[0]["colour"] == 0
+    assert G.nodes[1]["colour"] == 2
+    assert G.nodes[2]["colour"] == 4