Solve slow maintanance of blossom list

python/mwmatching.py

@@ -5,6 +5,7 @@ Algorithm for finding a maximum weight matching in general graphs.
 from __future__ import annotations
 
 import sys
+import itertools
 import math
 from collections.abc import Sequence
 from typing import NamedTuple, Optional
@@ -404,7 +405,7 @@ class _Blossom:
         self.vertex_dual_offset: float = 0
 
         # "marker" is a temporary variable used to discover common
-        # ancestors in the blossom tree. It is normally False, except
+        # ancestors in the alternating tree. It is normally False, except
         # when used by "trace_alternating_paths()".
         self.marker: bool = False
 
@@ -549,7 +550,7 @@ class _MatchingContext:
         # the course of the algorithm.
         #
         # Initially there are no non-trivial blossoms.
-        self.nontrivial_blossom: list[_NonTrivialBlossom] = []
+        self.nontrivial_blossom: set[_NonTrivialBlossom] = set()
 
         # "vertex_set_node[x]" represents the vertex "x" inside the
         # union-find datastructure of its top-level blossom.
@@ -616,12 +617,8 @@ class _MatchingContext:
     def __del__(self) -> None:
         """Delete reference cycles during cleanup of the matching context."""
 
-        for blossom in self.trivial_blossom:
+        for blossom in itertools.chain(self.trivial_blossom,
-            blossom.vertex_set.clear()
+                                       self.nontrivial_blossom):
-            del blossom.vertex_set
-            blossom.tree_blossoms = None
-
-        for blossom in self.nontrivial_blossom:
             blossom.vertex_set.clear()
             del blossom.vertex_set
             blossom.tree_blossoms = None
@@ -774,10 +771,6 @@ class _MatchingContext:
             # T-vertices.
             self.vertex_sedge_queue[x].clear()
             for e in self.graph.adjacent_edges[x]:
-                # TODO -- Postpone this cleanup step to the scanning of
-                #         S-vertex incident edges.
-                #         It will be more efficient, and also simpler to
-                #         reason about final time complexity.
                 self.vertex_sedge_node[e] = None
             self.vertex_set_node[x].set_prio(math.inf)
 
@@ -925,7 +918,8 @@ class _MatchingContext:
 
     def _check_alternating_tree_consistency(self) -> None:
         """TODO -- remove this function, only for debugging"""
-        for blossom in self.trivial_blossom + self.nontrivial_blossom:
+        for blossom in itertools.chain(self.trivial_blossom,
+                                       self.nontrivial_blossom):
             if (blossom.parent is None) and (blossom.label != _LABEL_NONE):
                 assert blossom.tree_blossoms is not None
                 assert blossom in blossom.tree_blossoms
@@ -938,38 +932,6 @@ class _MatchingContext:
                 assert blossom.tree_edge is None
                 assert blossom.tree_blossoms is None
 
-    def reset_stage(self) -> None:
-        """Reset data which are only valid during a stage.
-
-        Marks all blossoms as unlabeled, applies delayed delta updates,
-        and resets tracking of least-slack edges.
-
-        This function takes time O((n + m) * log(n)).
-        """
-
-        assert not self.scan_queue
-
-        for blossom in self.trivial_blossom + self.nontrivial_blossom:
-
-            # Remove blossom label.
-            if (blossom.parent is None) and (blossom.label != _LABEL_NONE):
-                self.reset_blossom_label(blossom)
-            assert blossom.label == _LABEL_NONE
-
-            # Remove blossom from alternating tree.
-            blossom.tree_edge = None
-            blossom.tree_blossoms = None
-
-            # Unwind lazy delta updates to vertex dual variables.
-            if blossom.vertex_dual_offset != 0:
-                for x in blossom.vertices():
-                    self.vertex_dual_2x[x] += blossom.vertex_dual_offset
-            blossom.vertex_dual_offset = 0
-
-        assert self.delta2_queue.empty()
-        assert self.delta3_queue.empty()
-        assert self.delta4_queue.empty()
-
     def remove_alternating_tree(
             self,
             tree_blossoms: set[_Blossom]
@@ -1123,7 +1085,7 @@ class _MatchingContext:
         tree_blossoms.add(blossom)
 
         # Insert into the blossom array.
-        self.nontrivial_blossom.append(blossom)
+        self.nontrivial_blossom.add(blossom)
 
         # Link the subblossoms to the their new parent.
         for sub in subblossoms:
@@ -1251,7 +1213,6 @@ class _MatchingContext:
             tree_blossoms.add(sub)
 
         # Delete the expanded blossom.
-        # TODO -- list manipulation is too slow
         self.nontrivial_blossom.remove(blossom)
 
     def expand_unlabeled_blossom(self, blossom: _NonTrivialBlossom) -> None:
@@ -1291,7 +1252,6 @@ class _MatchingContext:
                 sub.delta2_node = self.delta2_queue.insert(prio, sub)
 
         # Delete the expanded blossom.
-        # TODO -- list manipulation is too slow
         self.nontrivial_blossom.remove(blossom)
 
     #
@@ -1747,7 +1707,7 @@ class _MatchingContext:
         # This loop runs through at most O(n) iterations per stage.
         while True:
 
-            # self._check_alternating_tree_consistency()  # TODO -- remove this
+            self._check_alternating_tree_consistency()  # TODO -- remove this
 
             # Consider the incident edges of newly labeled S-vertices.
             self.substage_scan()
@@ -1807,10 +1767,32 @@ class _MatchingContext:
         Also resets tracking of least-slack edges.
 
         This function takes time O(n * log(n)).
-        It is called only once, at the end of the algorithm.
+        It is called once, at the end of the algorithm.
         """
-        # TODO -- move that function in here
+
-        self.reset_stage()
+        assert not self.scan_queue
+
+        for blossom in itertools.chain(self.trivial_blossom,
+                                       self.nontrivial_blossom):
+
+            # Remove blossom label.
+            if (blossom.parent is None) and (blossom.label != _LABEL_NONE):
+                self.reset_blossom_label(blossom)
+            assert blossom.label == _LABEL_NONE
+
+            # Remove blossom from alternating tree.
+            blossom.tree_edge = None
+            blossom.tree_blossoms = None
+
+            # Unwind lazy delta updates to vertex dual variables.
+            if blossom.vertex_dual_offset != 0:
+                for x in blossom.vertices():
+                    self.vertex_dual_2x[x] += blossom.vertex_dual_offset
+            blossom.vertex_dual_offset = 0
+
+        assert self.delta2_queue.empty()
+        assert self.delta3_queue.empty()
+        assert self.delta4_queue.empty()
 
 
 def _verify_blossom_edges(