Python NetworkX — set node color automatically based on a list of values

I generated a graph with networkx

import networkx as nx  
s = 5
G = nx.grid_graph(dim=[s,s])
nodes = list(G.nodes)
edges = list(G.edges)
p = []
for i in range(0, s):
    for j in range(0, s):
        p.append([i,j])
for i in range(0, len(nodes)):
    G.nodes[nodes[i]]['pos'] = p[i]

pos = {}
for i in range(0, len(nodes)):
        pos[nodes[i]] = p[i]

nx.draw(G, pos)

Now I would like to assign a value to each node between 0 and 4

from random import randint
val = []
for i in range(0, len(G.nodes())):
    val.append(randint(0,4)) 

And I would like to assign the color to each node base on the val list and plot something like shown here

Answer

To set a node property, you can use:

nx.set_node_attributes(G, val, 'val')

Networkx draw calls down to draw_networkx_nodes, this takes a cmap and color list, so all you would have to do would be something like:

nx.draw(G, pos, node_color = nx.get_node_attributes(G,'val'), vmin=0, vmax=4, cmap = plt.cm.get_cmap('rainbow'))