Add legend to NetworkX graph (using matplotlib) based on node attribute-CodePudding

I'm having trouble adding a legend to a NetworkX graph (have checked for previous questions but the q's & a's so far have been specific to only semi-related issues).

I'd like to be able to map my network and color the nodes according to node attributes (in the example I'm using here the attribute is 'sex' (m/f) but I have also attributes which have more categories or which are numerical). I'd also like a legend to show which color represents which category (in the case of categorical attributes).

So far, I've been able to color the nodes according to attribute but have been unable to add a legend to the graph. I've tried numerous things but nothing seems to be working atm!. Does anyone have any suggestions of how I can add a legend?

This is my code so far:

nodes_colors = []
for node in property_crime.nodes:
   if property_crime.nodes[node]["sex"] == "f":
       nodes_colors.append('blue')
   else:
       nodes_colors.append('orange')

nx.draw(property_crime, node_color=nodes_colors, node_size=15)
plt.show()

This is the image which is returned: enter image description here

Note that 'property_crime' is a graph object. The graph is undirected and is not weighted.

Thanks for any help!

CodePudding user response：

Try adding these lines before plt.show():

import matplotlib.pyplot as plt
from matplotlib.lines import Line2D

legend_elements = [Line2D([0], [0], marker='o', color='blue', label='Female', lw=0,
                          markerfacecolor='blue', markersize=10),
                   Line2D([0], [0], marker='o', color='orange', label='Male', lw=0,
                          markerfacecolor='orange', markersize=10)]
ax = plt.gca()
ax.legend(handles=legend_elements, loc='upper right')