diff --git a/python/mwmatching.py b/python/mwmatching.py
index f7f61ec..0116dc8 100644
--- a/python/mwmatching.py
+++ b/python/mwmatching.py
@@ -382,6 +382,12 @@ class _Blossom:
         # "tree_edge = None" if the blossom is the root of an alternating tree.
         self.tree_edge: Optional[tuple[int, int]] = None
 
+        # For a labeled top-level blossom,
+        # "alternating_tree" is the set of all top-level blossoms that belong
+        # to the same alternating tree. The same set instance is shared by
+        # all top-level blossoms in the tree.
+        self.tree_blossoms: "Optional[set[_Blossom]]" = None
+
         # Each top-level blossom maintains a union-find datastructure
         # containing all vertices in the blossom.
         self.vertex_set: "UnionFindQueue[_Blossom, int]"
@@ -607,10 +613,12 @@ class _MatchingContext:
         for blossom in self.trivial_blossom:
             blossom.vertex_set.clear()
             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
 
     def edge_pseudo_slack_2x(self, e: int) -> float:
         """Return 2 times the pseudo-slack of the specified edge.
@@ -860,6 +868,19 @@ class _MatchingContext:
 
         assert not self.scan_queue
 
+        # Check consistency of alternating tree.
+        for blossom in 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
+                if blossom.tree_edge is not None:
+                    bx = self.vertex_set_node[blossom.tree_edge[0]].find()
+                    by = self.vertex_set_node[blossom.tree_edge[1]].find()
+                    assert bx.tree_blossoms is blossom.tree_blossoms
+                    assert by.tree_blossoms is blossom.tree_blossoms
+            else:
+                assert blossom.tree_blossoms is None
+
         # Remove blossom labels and unwind lazy dual updates.
         for blossom in self.trivial_blossom + self.nontrivial_blossom:
             if blossom.parent is None:
@@ -868,6 +889,7 @@ class _MatchingContext:
                     blossom.delta4_node = None
             assert blossom.label == _LABEL_NONE
             blossom.tree_edge = None
+            blossom.tree_blossoms = None
 
         # Reset least-slack edge tracking.
         self.lset_reset()
@@ -944,7 +966,8 @@ class _MatchingContext:
         # If we found a common ancestor, trim the paths so they end there.
         if first_common is not None:
             assert self.vertex_set_node[xedges[-1][0]].find() is first_common
-            while self.vertex_set_node[yedges[-1][0]].find() is not first_common:
+            while (self.vertex_set_node[yedges[-1][0]].find()
+                   is not first_common):
                 yedges.pop()
 
         # Fuse the two paths.
@@ -991,12 +1014,16 @@ class _MatchingContext:
 
         # Remove blossom labels.
         # Mark vertices inside former T-blossoms as S-vertices.
+        tree_blossoms = subblossoms[0].tree_blossoms
+        assert tree_blossoms is not None
         for sub in subblossoms:
             if sub.label == _LABEL_S:
                 self.remove_blossom_label_s(sub)
             elif sub.label == _LABEL_T:
                 self.remove_blossom_label_t(sub)
                 self.assign_vertex_label_s(sub)
+            sub.tree_blossoms = None
+            tree_blossoms.remove(sub)
 
         # Create the new blossom object.
         blossom = _NonTrivialBlossom(subblossoms, path.edges)
@@ -1004,6 +1031,8 @@ class _MatchingContext:
         # Assign label S to the new blossom and link it to the tree.
         self.assign_blossom_label_s(blossom)
         blossom.tree_edge = subblossoms[0].tree_edge
+        blossom.tree_blossoms = tree_blossoms
+        tree_blossoms.add(blossom)
 
         # Insert into the blossom array.
         self.nontrivial_blossom.append(blossom)
@@ -1057,6 +1086,11 @@ class _MatchingContext:
         self.delta4_queue.delete(blossom.delta4_node)
         blossom.delta4_node = None
 
+        # Remove blossom from its alternating tree.
+        tree_blossoms = blossom.tree_blossoms
+        assert tree_blossoms is not None
+        tree_blossoms.remove(blossom)
+
         # Split union-find structure.
         blossom.vertex_set.split()
 
@@ -1094,6 +1128,8 @@ class _MatchingContext:
         # Assign label T to that sub-blossom.
         self.assign_blossom_label_t(sub)
         sub.tree_edge = blossom.tree_edge
+        sub.tree_blossoms = tree_blossoms
+        tree_blossoms.add(sub)
 
         # Walk through the expanded blossom from "sub" to the base vertex.
         # Assign alternating S and T labels to the sub-blossoms and attach
@@ -1118,6 +1154,8 @@ class _MatchingContext:
             sub = path_nodes[p+2]
             self.assign_blossom_label_t(sub)
             sub.tree_edge = path_edges[p+1]
+            sub.tree_blossoms = tree_blossoms
+            tree_blossoms.add(sub)
 
         # Delete the expanded blossom.
         # TODO -- list manipulation is too slow
@@ -1336,6 +1374,7 @@ class _MatchingContext:
 
             # Mark the blossom as root of an alternating tree.
             bx.tree_edge = None
+            bx.tree_blossoms = {bx}
 
         else:
             # Vertex "x" is matched to T-vertex "y".
@@ -1345,6 +1384,10 @@ class _MatchingContext:
             # Attach the blossom that contains "x" to the alternating tree.
             bx.tree_edge = (y, x)
 
+            bx.tree_blossoms = by.tree_blossoms
+            assert bx.tree_blossoms is not None
+            bx.tree_blossoms.add(bx)
+
     def extend_tree_t(self, x: int, y: int) -> None:
         """Assign label T to the unlabeled blossom that contains vertex "y".
 
@@ -1356,9 +1399,10 @@ class _MatchingContext:
          - "y" is an unlabeled, matched vertex.
          - There is a tight edge between vertices "x" and "y".
         """
-        assert self.vertex_set_node[x].find().label == _LABEL_S
 
+        bx = self.vertex_set_node[x].find()
         by = self.vertex_set_node[y].find()
+        assert bx.label == _LABEL_S
 
         # Expand zero-dual blossoms before assigning label T.
         while isinstance(by, _NonTrivialBlossom) and (by.dual_var == 0):
@@ -1368,6 +1412,9 @@ class _MatchingContext:
         # Assign label T to the unlabeled blossom.
         self.assign_blossom_label_t(by)
         by.tree_edge = (x, y)
+        by.tree_blossoms = bx.tree_blossoms
+        assert by.tree_blossoms is not None
+        by.tree_blossoms.add(by)
 
         # Assign label S to the blossom that is mated to the T-blossom.
         z = self.vertex_mate[by.base_vertex]
@@ -1427,7 +1474,7 @@ class _MatchingContext:
             # Scan the edges that are incident on "x".
             # This loop runs through O(m) iterations per stage.
             for e in adjacent_edges[x]:
-                (p, q, w) = edges[e]
+                (p, q, _w) = edges[e]
                 y = p if p != x else q
 
                 # Consider the edge between vertices "x" and "y".
@@ -1600,7 +1647,7 @@ class _MatchingContext:
 
             # Calculate delta step in the dual LPP problem.
             (delta_type, delta_2x, delta_edge, delta_blossom
-                ) = self.substage_calc_dual_delta()
+             ) = self.substage_calc_dual_delta()
 
             # Update the running sum of delta steps.
             # This implicitly updates the dual variables as needed, because