From 271748700f127bdd63e51f8e9da521de4395149a Mon Sep 17 00:00:00 2001
From: Manuel Raimann <raimannma@outlook.de>
Date: Thu, 12 Feb 2026 14:13:34 +0100
Subject: [PATCH] perf(tpe): use quickselect instead of full sort in
 split_trials

---
 src/sampler/tpe/sampler.rs | 31 ++++++++++++++-----------------
 1 file changed, 14 insertions(+), 17 deletions(-)

diff --git a/src/sampler/tpe/sampler.rs b/src/sampler/tpe/sampler.rs
index 16ff43b..47d2d09 100644
--- a/src/sampler/tpe/sampler.rs
+++ b/src/sampler/tpe/sampler.rs
@@ -288,15 +288,6 @@ impl TpeSampler {
             return (vec![], vec![]);
         }
 
-        // Sort trials by value (ascending for minimization)
-        let mut sorted_indices: Vec<usize> = (0..history.len()).collect();
-        sorted_indices.sort_by(|&a, &b| {
-            history[a]
-                .value
-                .partial_cmp(&history[b].value)
-                .unwrap_or(core::cmp::Ordering::Equal)
-        });
-
         // Compute gamma using the strategy and clamp to valid range
         let gamma = self
             .gamma_strategy
@@ -309,14 +300,20 @@ impl TpeSampler {
             .max(1)
             .min(history.len() - 1);
 
-        let good: Vec<_> = sorted_indices[..n_good]
-            .iter()
-            .map(|&i| &history[i])
-            .collect();
-        let bad: Vec<_> = sorted_indices[n_good..]
-            .iter()
-            .map(|&i| &history[i])
-            .collect();
+        // Use quickselect (O(n)) to partition indices instead of full sort (O(n log n)).
+        // We only need to know which trials are in the top gamma-quantile, not their order.
+        let mut indices: Vec<usize> = (0..history.len()).collect();
+        if n_good > 0 {
+            indices.select_nth_unstable_by(n_good - 1, |&a, &b| {
+                history[a]
+                    .value
+                    .partial_cmp(&history[b].value)
+                    .unwrap_or(core::cmp::Ordering::Equal)
+            });
+        }
+
+        let good: Vec<_> = indices[..n_good].iter().map(|&i| &history[i]).collect();
+        let bad: Vec<_> = indices[n_good..].iter().map(|&i| &history[i]).collect();
 
         (good, bad)
     }