From deab8938e502ea8d956aa3d2ed5500baa5838e36 Mon Sep 17 00:00:00 2001
From: Chris Boseak <cboseak@gmail.com>
Date: Fri, 4 Sep 2026 07:02:20 -0500
Subject: [PATCH] FIX Restore UnionFind path compression

---
 asv_benchmarks/benchmarks/cluster.py          | 27 ++++++++++++++
 .../sklearn.cluster/34872.efficiency.rst      |  6 ++++
 sklearn/cluster/_hierarchical_fast.pyx        | 35 ++++++++++++++-----
 sklearn/cluster/tests/test_hierarchical.py    | 31 ++++++++++++++++
 4 files changed, 90 insertions(+), 9 deletions(-)
 create mode 100644 doc/whats_new/upcoming_changes/sklearn.cluster/34872.efficiency.rst

diff --git a/asv_benchmarks/benchmarks/cluster.py b/asv_benchmarks/benchmarks/cluster.py
index 457a15dd938e9..8986870663431 100644
--- a/asv_benchmarks/benchmarks/cluster.py
+++ b/asv_benchmarks/benchmarks/cluster.py
@@ -1,3 +1,6 @@
+import numpy as np
+from sklearn.cluster._hierarchical_fast import _single_linkage_label
+
 from sklearn.cluster import KMeans, MiniBatchKMeans
 
 from .common import Benchmark, Estimator, Predictor, Transformer
@@ -5,6 +8,30 @@
 from .utils import neg_mean_inertia
 
 
+class SingleLinkageLabelBenchmark(Benchmark):
+    """Benchmark conversion of an adversarial MST to single linkage."""
+
+    param_names = ["n_samples"]
+    params = ([2_000, 16_000],)
+
+    def setup(self, n_samples):
+        half_n_samples = n_samples // 2
+        self.mst = np.empty((n_samples - 1, 3), dtype=np.float64)
+
+        indices = np.arange(half_n_samples - 1)
+        self.mst[: half_n_samples - 1, 0] = indices
+        self.mst[: half_n_samples - 1, 1] = indices + 1
+
+        indices = np.arange(half_n_samples)
+        self.mst[half_n_samples - 1 :, 0] = indices
+        self.mst[half_n_samples - 1 :, 1] = half_n_samples + indices
+
+        self.mst[:, 2] = np.arange(n_samples - 1)
+
+    def time_single_linkage_label(self, n_samples):
+        _single_linkage_label(self.mst)
+
+
 class KMeansBenchmark(Predictor, Transformer, Estimator, Benchmark):
     """
     Benchmarks for KMeans.
diff --git a/doc/whats_new/upcoming_changes/sklearn.cluster/34872.efficiency.rst b/doc/whats_new/upcoming_changes/sklearn.cluster/34872.efficiency.rst
new file mode 100644
index 0000000000000..1f29c39ed6d4d
--- /dev/null
+++ b/doc/whats_new/upcoming_changes/sklearn.cluster/34872.efficiency.rst
@@ -0,0 +1,6 @@
+- :class:`cluster.AgglomerativeClustering` and
+  :class:`cluster.FeatureAgglomeration` with `linkage="single"`, as well as
+  :class:`cluster.HDBSCAN`, are now faster for inputs that produce deep
+  union-find trees. Restored path compression avoids quadratic time while
+  constructing the single-linkage hierarchy.
+  By :user:`Chris Boseak <cboseak>`.
diff --git a/sklearn/cluster/_hierarchical_fast.pyx b/sklearn/cluster/_hierarchical_fast.pyx
index 8d7c363daef37..0608e58de8a46 100644
--- a/sklearn/cluster/_hierarchical_fast.pyx
+++ b/sklearn/cluster/_hierarchical_fast.pyx
@@ -336,17 +336,34 @@ cdef class UnionFind(object):
         self.next_label += 1
         return
 
-    @cython.wraparound(True)
     cdef intp_t fast_find(self, intp_t n) noexcept:
-        cdef intp_t p
-        p = n
+        cdef intp_t root = n
+        cdef intp_t next_parent
+
         # find the highest node in the linkage graph so far
-        while self.parent[n] != -1:
-            n = self.parent[n]
-        # provide a shortcut up to the highest node
-        while self.parent[p] != n:
-            p, self.parent[p] = self.parent[p], n
-        return n
+        while self.parent[root] != -1:
+            root = self.parent[root]
+        # compress the path to the highest node
+        while n != root and self.parent[n] != root:
+            next_parent = self.parent[n]
+            self.parent[n] = root
+            n = next_parent
+        return root
+
+
+cdef class PytestUnionFind(UnionFind):
+    """Used for testing only."""
+
+    def py_union(self, intp_t m, intp_t n):
+        cdef intp_t new_label = self.next_label
+        self.union(m, n)
+        return new_label
+
+    def py_fast_find(self, intp_t n):
+        return self.fast_find(n)
+
+    def py_get_parent(self):
+        return np.asarray(self.parent).copy()
 
 
 def _single_linkage_label(const float64_t[:, :] L):
diff --git a/sklearn/cluster/tests/test_hierarchical.py b/sklearn/cluster/tests/test_hierarchical.py
index 66488c6338f07..e4746aacc043a 100644
--- a/sklearn/cluster/tests/test_hierarchical.py
+++ b/sklearn/cluster/tests/test_hierarchical.py
@@ -24,6 +24,7 @@
     linkage_tree,
 )
 from sklearn.cluster._hierarchical_fast import (
+    PytestUnionFind,
     average_merge,
     max_merge,
     mst_linkage_core,
@@ -387,6 +388,36 @@ def test_vector_scikit_single_vs_scipy_single(global_random_seed):
     assess_same_labelling(cut, cut_scipy)
 
 
+def test_union_find_fast_find_compresses_path():
+    """Check that fast_find compresses the queried path to its root."""
+    union_find = PytestUnionFind(5)
+
+    node_5 = union_find.py_union(0, 1)
+    node_6 = union_find.py_union(node_5, 2)
+    root = union_find.py_union(node_6, 3)
+
+    assert union_find.py_fast_find(0) == root
+
+    parent = union_find.py_get_parent()
+    assert_array_equal(
+        parent[[0, node_5, node_6]],
+        np.full(3, root, dtype=np.intp),
+    )
+    assert parent[root] == -1
+
+
+def test_union_find_fast_find_on_root_is_noop():
+    """Check that fast_find does not mutate state when called on a root.
+
+    Non-regression test for issue #34626.
+    """
+    union_find = PytestUnionFind(3)
+    parent_before = union_find.py_get_parent()
+
+    assert union_find.py_fast_find(0) == 0
+    assert_array_equal(union_find.py_get_parent(), parent_before)
+
+
 @pytest.mark.parametrize("metric_param_grid", METRICS_DEFAULT_PARAMS)
 def test_mst_linkage_core_memory_mapped(metric_param_grid):
     """The MST-LINKAGE-CORE algorithm must work on mem-mapped dataset.