Implemented multipartite_layout (#3815)

* Implemented multipartite_layout

* Added test.

* Remove left over `max_layer` in multipartite_layout

* Multipartite Layout Example

* Update examples/drawing/plot_multipartite_graph.py

Co-authored-by: Jarrod Millman <jarrod.millman@gmail.com>

1 files changed, 39 insertions, 0 deletions

diff --git a/examples/drawing/plot_multipartite_graph.py b/examples/drawing/plot_multipartite_graph.py
new file mode 100644
index 00000000..87da9616
--- /dev/null
+++ b/examples/drawing/plot_multipartite_graph.py
@@ -0,0 +1,39 @@
+"""
+===================
+Multipartite Layout
+===================
+"""
+
+import itertools
+import matplotlib.pyplot as plt
+import networkx as nx
+
+from networkx.utils import pairwise
+
+subset_sizes = [5,5,4,3,2,4,4,3]
+subset_color = ['gold',
+                'violet',
+                'violet',
+                'violet',
+                'violet',
+                'limegreen',
+                'limegreen',
+                'darkorange']
+
+def multilayered_graph(*subset_sizes):
+    extents = pairwise(itertools.accumulate((0,) + subset_sizes))
+    layers = [range(start, end) for start, end in extents]
+    G = nx.Graph()
+    for (i, layer) in enumerate(layers):
+        G.add_nodes_from(layer, layer=i)
+    for layer1, layer2 in pairwise(layers):
+        G.add_edges_from(itertools.product(layer1, layer2))
+    return G
+
+G = multilayered_graph(*subset_sizes)
+color = [subset_color[data['layer']]  for v,data in G.nodes(data=True)]
+pos = nx.multipartite_layout(G,subset_key='layer')
+plt.figure(figsize=(8, 8))
+nx.draw(G, pos, node_color=color, with_labels=False)
+plt.axis("equal")
+plt.show()