diff --git a/listcolouring/__init__.py b/listcolouring/__init__.py index f5ff9e1..56e91b6 100644 --- a/listcolouring/__init__.py +++ b/listcolouring/__init__.py @@ -51,12 +51,17 @@ def first_permissible_or_none(G, u, v): choice = None return(choice) -def first_permissible_or_none_node(G, u): +def first_permissible_or_none_node(G, u, edges = False): """ Returns the first element of A = P - X if A is non-empty otherwise returns None. - X is the list of colours on neighbours of u. - P is the list of permissible colours for node u. + X is the list of colours on neighbours of u if edges = False and the union of + the list of colours on neighbours of u and the list of colours on edges incident + with u if edges = True. + P is the list of permissible colours for node u. """ - X = colours_on_neighbours(G, u) + if edges: + X = colours_on_neighbours(G, u).union(colours_incident_with(G, u)) + else: + X = colours_on_neighbours(G, u) P = set(G.nodes[u]["permissible"]) choices = P - X if(len(choices) > 0): diff --git a/tests/test_first-permissible-or-none-node.py b/tests/test_first-permissible-or-none-node.py index 5eac92b..39736bb 100644 --- a/tests/test_first-permissible-or-none-node.py +++ b/tests/test_first-permissible-or-none-node.py @@ -11,3 +11,7 @@ def test_first_permissible_or_none_node(): assert first_permissible_or_none_node(G, 0) == 0 assert first_permissible_or_none_node(G, 1) == 1 assert first_permissible_or_none_node(G, 2) == 2 + nx.set_edge_attributes(G, 2, "colour") + assert first_permissible_or_none_node(G, 0, edges = True) == 0 + assert first_permissible_or_none_node(G, 1, edges = True) == 1 + assert first_permissible_or_none_node(G, 2, edges = True) == 3