Implement colours_on_neighbours. (#33)

listcolouring/__init__.py

@@ -29,6 +29,10 @@ def colours_incident_with(G, u):
     """The list of colours on edges incident with vertex u in graph G."""
     return(set([G[u][v]["colour"] for v in nx.neighbors(G, u)]))
 
+def colours_on_neighbours(G, n):
+    """The set of all colours on neighbours of a node n in a graph G."""
+    return(set([nx.get_node_attributes(G, "colour")[m] for m in G.neighbors(n)]))
+
 def first_permissible_or_none(G, u, v):
     """
     Returns the first element of A if A is non-empty otherwise returns None.

tests/test_colours_on_neighbours.py

@@ -0,0 +1,13 @@
+import networkx as nx
+
+import listcolouring
+from listcolouring import colours_on_neighbours
+
+def test_colours_on_neighbours():
+    G = nx.complete_graph(3)
+    node_colours_dict = {0: 3, 1: 2, 2:1}
+    nx.set_node_attributes(G, node_colours_dict, "colour")
+    assert colours_on_neighbours(G, 0) == set([2, 1])
+    assert colours_on_neighbours(G, 1) == set([3, 1])
+    assert colours_on_neighbours(G, 2) == set([3, 2])
+