From 6436a1d8f2bda07ce640714ad806498b21e46212 Mon Sep 17 00:00:00 2001 From: Matthew Henderson Date: Wed, 6 Nov 2024 09:20:56 +0000 Subject: [PATCH] Implement colours_on_neighbours. (#33) --- listcolouring/__init__.py | 4 ++++ tests/test_colours_on_neighbours.py | 13 +++++++++++++ 2 files changed, 17 insertions(+) create mode 100644 tests/test_colours_on_neighbours.py diff --git a/listcolouring/__init__.py b/listcolouring/__init__.py index 8307339..5b00587 100644 --- a/listcolouring/__init__.py +++ b/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. diff --git a/tests/test_colours_on_neighbours.py b/tests/test_colours_on_neighbours.py new file mode 100644 index 0000000..2e0d827 --- /dev/null +++ b/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]) +