Clean up explicit-stack recursion

python/max_weight_matching.py

@@ -1137,37 +1137,15 @@ class _MatchingContext:
         self.blossom[b] = None
         self.unused_blossoms.append(b)
 
-    def expand_zero_dual_blossoms(self) -> None:
+    def expand_blossom_rec(self, b: int, stack: list[int]) -> None:
-        """Expand all blossoms with zero dual variable (recursively).
+        """Expand blossom "b" and recursively expand any sub-blossoms that
+        have dual variable zero.
 
-        Note that this function runs at the end of a stage.
+        Use the stack object instead of making direct recursive calls.
-        Blossoms are not labeled. Least-slack edges are not tracked.
-
-        This function takes time O(n).
         """
 
         num_vertex = self.graph.num_vertex
 
-# TODO : clean up explicit stack
-
-        # Find top-level blossoms with zero slack.
-        stack: list[int] = []
-        for b in range(num_vertex, 2 * num_vertex):
-            blossom = self.blossom[b]
-            if (blossom is not None) and (self.blossom_parent[b] == -1):
-                # We typically expand only S-blossoms that were created after
-                # the most recent delta step. Those blossoms have _exactly_
-                # zero dual. So this comparison is reliable, even in case
-                # of floating point edge weights.
-                if blossom.dual_var == 0:
-                    stack.append(b)
-
-        # Use an explicit stack to avoid deep recursion.
-        while stack:
-            b = stack.pop()
-
-            # Expand blossom "b".
-
         blossom = self.blossom[b]
         assert blossom is not None
         assert self.blossom_parent[b] == -1
@@ -1197,29 +1175,54 @@ class _MatchingContext:
         self.blossom[b] = None
         self.unused_blossoms.append(b)
 
-    def augment_blossom(self, b: int, sub: int) -> None:
+    def expand_zero_dual_blossoms(self) -> None:
-        """Augment along an alternating path through blossom "b",
+        """Expand all blossoms with zero dual variable (recursively).
-        from sub-blossom "sub" to the base vertex of the blossom.
+
+        Note that this function runs at the end of a stage.
+        Blossoms are not labeled. Least-slack edges are not tracked.
 
         This function takes time O(n).
         """
+
         num_vertex = self.graph.num_vertex
 
-# TODO : cleanup explicit stack
-
         # Use an explicit stack to avoid deep recursion.
-        stack = [(b, sub)]
+        # The stack contains a list of blossoms to be expanded.
+        stack: list[int] = []
 
+        # Find top-level blossoms with zero slack.
+        for b in range(num_vertex, 2 * num_vertex):
+            blossom = self.blossom[b]
+            if (blossom is not None) and (self.blossom_parent[b] == -1):
+                # We typically expand only S-blossoms that were created after
+                # the most recent delta step. Those blossoms have _exactly_
+                # zero dual. So this comparison is reliable, even in case
+                # of floating point edge weights.
+                if blossom.dual_var == 0:
+                    stack.append(b)
+
+        # Expand blossoms.
         while stack:
-            (top_blossom, sub) = stack.pop()
+            b = stack.pop()
-            b = self.blossom_parent[sub]
+            self.expand_blossom_rec(b, stack)
 
-            if b != top_blossom:
+    def augment_blossom_rec(
-                # Set up to continue augmenting through the parent of "b".
+            self,
-                stack.append((top_blossom, b))
+            b: int,
+            sub: int,
+            stack: list[tuple[int, int]]
+            ) -> None:
+        """Augment along an alternating path through blossom "b",
+        from sub-blossom "sub" to the base vertex of the blossom.
 
-            # Augment blossom "b" from subblossom "sub" to the base of the
+        Modify the blossom to reflect that sub-blossom "sub" contains
-            # blossom. Afterwards, "sub" contains the new base vertex.
+        the base vertex after augmenting.
+
+        Recursively augment any sub-blossoms on the alternating path,
+        except for sub-blossom "sub" which has already been augmented.
+        Use the stack object instead of making direct recursive calls.
+        """
+        num_vertex = self.graph.num_vertex
 
         blossom = self.blossom[b]
         assert blossom is not None
@@ -1268,6 +1271,36 @@ class _MatchingContext:
         else:
             blossom.base_vertex = self.get_blossom(sub).base_vertex
 
+    def augment_blossom(self, b: int, sub: int) -> None:
+        """Augment along an alternating path through blossom "b",
+        from sub-blossom "sub" to the base vertex of the blossom.
+
+        This function takes time O(n).
+        """
+
+        # Use an explicit stack to avoid deep recursion.
+        stack = [(b, sub)]
+
+        while stack:
+            (outer_blossom, sub) = stack.pop()
+            b = self.blossom_parent[sub]
+
+            if b != outer_blossom:
+                # Sub-blossom "sub" is an indirect (nested) child of
+                # the blossom we are supposed to be augmenting.
+                #
+                # Blossom "b" is the direct parent of "sub".
+                # Let's first augment "b" from "sub" to its base vertex.
+                # Then continue bay augmenting the parent of "b" from
+                # "b" to its base vertex, and so on until we get to the
+                # outer blossom.
+                #
+                # Set up to continue augmenting through the parent of "b".
+                stack.append((outer_blossom, b))
+
+            # Augment blossom "b" from "sub" to the base vertex of "b".
+            self.augment_blossom_rec(b, sub, stack)
+
     def augment_matching(self, path: _AlternatingPath) -> None:
         """Augment the matching through the specified augmenting path.
 
@@ -1440,7 +1473,7 @@ class _MatchingContext:
             # This loop runs through O(m) iterations per stage.
             for e in adjacent_edges[x]:
                 (p, q, _w) = edges[e]
-                y = p if p != x else q  # TODO : consider abstracting this
+                y = p if p != x else q
 
                 # Consider the edge between vertices "x" and "y".
                 # Try to pull this edge into an alternating tree.