modularity_max: breaking the loop when given community size is reached (#4950)

* modularity_max: allow input of desired number of communities

* import warnings

* format

* format

* improvements according to discussion

* try to manually merge main + resolve conflicts

* add test for n_communities parameter using circular ladder graph

* style of test

2 files changed, 34 insertions, 4 deletions

diff --git a/networkx/algorithms/community/modularity_max.py b/networkx/algorithms/community/modularity_max.py
index 2e608a7e..debf5fb3 100644
--- a/networkx/algorithms/community/modularity_max.py
+++ b/networkx/algorithms/community/modularity_max.py
@@ -11,7 +11,7 @@ __all__ = [
 ]
 
 
-def greedy_modularity_communities(G, weight=None, resolution=1):
+def greedy_modularity_communities(G, weight=None, resolution=1, n_communities=1):
     r"""Find communities in G using greedy modularity maximization.
 
     This function uses Clauset-Newman-Moore greedy modularity maximization [2]_.
@@ -19,7 +19,7 @@ def greedy_modularity_communities(G, weight=None, resolution=1):
 
     Greedy modularity maximization begins with each node in its own community
     and joins the pair of communities that most increases modularity until no
-    such pair exists.
+    such pair exists or until number of communities `n_communities` is reached.
 
     This function maximizes the generalized modularity, where `resolution`
     is the resolution parameter, often expressed as $\gamma$.
@@ -38,6 +38,14 @@ def greedy_modularity_communities(G, weight=None, resolution=1):
         If resolution is less than 1, modularity favors larger communities.
         Greater than 1 favors smaller communities.
 
+    n_communities: int
+        Desired number of communities: the community merging process is
+        terminated once this number of communities is reached, or until
+        modularity can not be further increased. Must be between 1 and the
+        total number of nodes in `G`. Default is ``1``, meaning the community
+        merging process continues until all nodes are in the same community
+        or until the best community structure is found.
+
     Returns
     -------
     list
@@ -72,6 +80,11 @@ def greedy_modularity_communities(G, weight=None, resolution=1):
     m = sum([d.get(weight, 1) for u, v, d in G.edges(data=True)])
     q0 = 1.0 / (2.0 * m)
 
+    if (n_communities < 1) or (n_communities > N):
+        raise ValueError(
+            f"n_communities must be between 1 and {len(G.nodes())}. Got {n_communities}"
+        )
+
     # Map node labels to contiguous integers
     label_for_node = {i: v for i, v in enumerate(G.nodes())}
     node_for_label = {label_for_node[i]: i for i in range(N)}
@@ -111,8 +124,9 @@ def greedy_modularity_communities(G, weight=None, resolution=1):
     ]
     H = MappedQueue([dq_heap[i].h[0] for i in range(N) if len(dq_heap[i]) > 0])
 
-    # Merge communities until we can't improve modularity
-    while len(H) > 1:
+    # Merge communities until we can't improve modularity or until desired number of
+    # communities (n_communities) is reached.
+    while len(H) > n_communities:
         # Find best merge
         # Remove from heap of row maxes
         # Ties will be broken by choosing the pair with lowest min community id
diff --git a/networkx/algorithms/community/tests/test_modularity_max.py b/networkx/algorithms/community/tests/test_modularity_max.py
index 11e6a00d..8a0cddbc 100644
--- a/networkx/algorithms/community/tests/test_modularity_max.py
+++ b/networkx/algorithms/community/tests/test_modularity_max.py
@@ -83,3 +83,19 @@ def test_resolution_parameter_impact():
     expected = [frozenset(range(8)), frozenset(range(8, 13))]
     assert greedy_modularity_communities(G, resolution=gamma) == expected
     assert naive_greedy_modularity_communities(G, resolution=gamma) == expected
+
+
+def test_n_communities_parameter():
+    G = nx.circular_ladder_graph(4)
+
+    # No aggregation:
+    expected = [{k} for k in range(8)]
+    assert greedy_modularity_communities(G, n_communities=8) == expected
+
+    # Aggregation to half order (number of nodes)
+    expected = [{k, k + 1} for k in range(0, 8, 2)]
+    assert greedy_modularity_communities(G, n_communities=4) == expected
+
+    # Default aggregation case (here, 2 communities emerge)
+    expected = [frozenset(range(0, 4)), frozenset(range(4, 8))]
+    assert greedy_modularity_communities(G, n_communities=1) == expected