So I'm using a networkX graph to represent some information. This information is represented by different object types (for example, ColorNode and ShapeNode).

Some of the processing that is done on this graph requires me to extract out a specific type of node. Every time I need to do this, I do something along the lines of the code below.

colornodes = []
for node in graph.nodes():
    if isinstance(node, ColorNode):

While this works, I feel like this is a situation that would arise often when working with graphs and I am re-inventing the wheel there. Essentially, I would like to know if there is a nicer way of doing this.

2 Answers

Kikohs On Best Solutions

Instead of defining your own type and always check with isinstance (which is painfully slow) I suggest another approach.

You can have a look at this answer for the classic node/edge filtering.

However I found another trick which may come handy for your specific case.

If you define an attribute that represents the node type, you can query nodes having that specific attribute using the builtin get_node_attributes function. The trick is that it only returns nodes that really define the attribute:

import networkx as nx

G = nx.complete_graph(10)
G.node[0]['ColorNode'] = True  # right-hand side value is irrelevant for the lookup
G.node[1]['ColorNode'] = True
G.node[2]['ShapeNode'] = True
G.node[3]['ShapeNode'] = True

shape_nodes = nx.get_node_attributes(G, 'ShapeNode').keys()
color_nodes = nx.get_node_attributes(G, 'ColorNode').keys()

print('Shape node ids: {}'.format(shape_nodes))
print('Color node ids: {}'.format(color_nodes))


Shape node ids: [2, 3]
Color node ids: [0, 1]

Of course if your graph is big or static, you should keep the id lists for fast querying!

sneha On

instead of graph.nodes(), you can use xpath.