taga
taga

Reputation: 3885

Setting the colour of Networkx nodes with Python dict

I want to set the colour for specific nodes in networkx lib with Python. These are my nodes:

node_list = ["A", "B", "C", "D", "E", "F"]

And these are the colours that I want:

A->red
B->red
C->red
D->green
E->green
F->blue

I do not know how to do this, because sometimes i have more than 200 different nodes so i am looking for some solution that will work with dict:

{"red":[A,B,C], "green": [D,E], "blue": [F]}

This is my code:

import networkx as nx

analysis_graph = nx.Graph()

node_list = ["A", "B", "C", "D", "E", "F"]

analysis_graph.add_nodes_from(node_list)
#print(analysis_graph.nodes())

nx.draw(analysis_graph, with_labels = True)
relation_list = [['A', 'B'],
                 ['A', 'C'],
                 ['B', 'D'],
                 ['C', 'E'],
                 ['D', 'F'],
                 ['E', 'D'],
                 ['C', 'E'],
                 ['B', 'D'],                 
                 ['C', 'F'],
                 ['A', 'E'],
                 ['B', 'C'],                 
                 ['B', 'F'],
                 ['D', 'A']]

analysis_graph = nx.from_edgelist(relation_list)
print(nx.info(analysis_graph))
nx.draw(analysis_graph, with_labels = True)

Upvotes: 0

Views: 203

Answers (1)

AKX
AKX

Reputation: 168824

Sure. According to the docs nx.draw supports a node_color argument. You'll just have to explicitly list the nodes in a proper order:

relation_list = [
    ["A", "B"],
    ["A", "C"],
    ["B", "D"],
    ["C", "E"],
    ["D", "F"],
    ["E", "D"],
    ["C", "E"],
    ["B", "D"],
    ["C", "F"],
    ["A", "E"],
    ["B", "C"],
    ["B", "F"],
    ["D", "A"],
]

colors = {"red": ["A", "B", "C"], "green": ["D", "E"], "blue": ["F"]}

analysis_graph = nx.from_edgelist(relation_list)


# Create a mapping of node -> color from `colors`:

node_color_map = {}
for color, nodes in colors.items():
    node_color_map.update(dict.fromkeys(nodes, color))

# Create a list of all the nodes to draw
node_list = sorted(nx.nodes(analysis_graph))
# ... and a list of colors in the proper order.
node_color_list = [node_color_map.get(node) for node in node_list]
# Draw the graph with the correct colors.
nx.draw(
    analysis_graph,
    with_labels=True,
    nodelist=node_list,
    node_color=node_color_list,
)

The output is e.g. enter image description here

Upvotes: 1

Related Questions