From a47182a8ff0d35f6fa600e928c3498a850eb2e8e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=BAlio=20Hoffimann?= <julio.hoffimann@gmail.com>
Date: Sun, 19 Jul 2026 12:56:57 -0300
Subject: [PATCH 1/3] Use KDTree beyond nmax in KMedoids when possible

---
 Project.toml               |  4 ++-
 src/TableTransforms.jl     |  1 +
 src/transforms/kmedoids.jl | 52 +++++++++++++++++++++++++-------------
 3 files changed, 38 insertions(+), 19 deletions(-)

diff --git a/Project.toml b/Project.toml
index f49e779f..e7203fbc 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "TableTransforms"
 uuid = "0d432bfd-3ee1-4ac1-886a-39f05cc69a3e"
-authors = ["Júlio Hoffimann <julio.hoffimann@gmail.com> and contributors"]
 version = "1.38.9"
+authors = ["Júlio Hoffimann <julio.hoffimann@gmail.com> and contributors"]
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
@@ -13,6 +13,7 @@ Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+NearestNeighbors = "b8a86587-4115-5ab1-83bc-aa920d37bbce"
 NelderMead = "2f6b4ddb-b4ff-44c0-b59b-2ab99302f970"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -34,6 +35,7 @@ Distributed = "1.10"
 Distributions = "0.25"
 InverseFunctions = "0.1"
 LinearAlgebra = "1.10"
+NearestNeighbors = "0.4"
 NelderMead = "0.4"
 PrettyTables = "3.0"
 Random = "1.10"
diff --git a/src/TableTransforms.jl b/src/TableTransforms.jl
index 25924584..16490383 100644
--- a/src/TableTransforms.jl
+++ b/src/TableTransforms.jl
@@ -26,6 +26,7 @@ using InverseFunctions: NoInverse, inverse as invfun
 using StatsBase: AbstractWeights, Weights, sample
 using Distributed: CachingPool, pmap, workers
 using SpectralIndices: indices, bands, compute
+using NearestNeighbors: KDTree, nn
 using NelderMead: optimise
 
 import Distributions: quantile, cdf
diff --git a/src/transforms/kmedoids.jl b/src/transforms/kmedoids.jl
index bbd48120..32a94a99 100644
--- a/src/transforms/kmedoids.jl
+++ b/src/transforms/kmedoids.jl
@@ -129,14 +129,9 @@ end
 
 function _updatelabels!(td, table, medoids, labels, dists)
   for (k, mₖ) in enumerate(medoids)
-    inds = 1:_nrows(table)
-
-    X = Tables.subset(table, inds, viewhint=true)
     μ = Tables.subset(table, [mₖ], viewhint=true)
-
-    δ = pairwise(td, X, μ)
-
-    @inbounds for i in inds
+    δ = pairwise(td, table, μ) |> vec
+    @inbounds for i in eachindex(δ)
       if δ[i] < dists[i]
         dists[i] = δ[i]
         labels[i] = k
@@ -152,20 +147,14 @@ function _updatemedoids!(td, table, medoids, labels)
     # if cluster is empty, then no medoid to update
     isempty(inds) && continue
 
-    X = Tables.subset(table, inds, viewhint=true)
-
-    j = _medoid(td, X)
+    # compute new medoid for cluster k
+    clust = Tables.subset(table, inds, viewhint=true)
+    _, j = findmin(sum, eachcol(pairwise(td, clust)))
 
     @inbounds medoids[k] = inds[j]
   end
 end
 
-function _medoid(td, table)
-  Δ = pairwise(td, table)
-  _, j = findmin(sum, eachcol(Δ))
-  j
-end
-
 function _interp(labels, inds, table, td)
   nobs = _nrows(table)
 
@@ -173,12 +162,39 @@ function _interp(labels, inds, table, td)
   ilabels[inds] .= labels
 
   X = Tables.subset(table, inds, viewhint=true)
+  s = _searcher(X, td)
+
   for i in setdiff(1:nobs, inds)
     x = Tables.subset(table, [i], viewhint=true)
-    δ = pairwise(td, X, x)
-    _, j = findmin(δ)
+    j = _search(s, x)
     ilabels[i] = labels[j]
   end
 
   ilabels
 end
+
+function _searcher(X, td)
+  # check if all variables are continuous
+  cols = Tables.columns(X)
+  vars = Tables.columnnames(cols)
+  allcont = all(vars) do var
+    x = Tables.getcolumn(cols, var)
+    elscitype(x) <: Continuous
+  end
+
+  # use KDTree if all variables are continuous
+  if allcont
+    data = Tables.matrix(X)
+    KDTree(transpose(data))
+  else
+    X, td
+  end
+end
+
+function _search(s::KDTree, x)
+  c = Tables.columns(x)
+  v = Tables.getcolumn(c, 1)
+  nn(s, v) |> first
+end
+
+_search((X, td), x) = pairwise(td, X, x) |> vec |> argmin

From 6a076057c6d0288486bd28cda5403f8c27b8289b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=BAlio=20Hoffimann?= <julio.hoffimann@gmail.com>
Date: Sun, 19 Jul 2026 13:13:04 -0300
Subject: [PATCH 2/3] Fix issue with column extraction

---
 src/transforms/kmedoids.jl | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/src/transforms/kmedoids.jl b/src/transforms/kmedoids.jl
index 32a94a99..247c4822 100644
--- a/src/transforms/kmedoids.jl
+++ b/src/transforms/kmedoids.jl
@@ -191,10 +191,6 @@ function _searcher(X, td)
   end
 end
 
-function _search(s::KDTree, x)
-  c = Tables.columns(x)
-  v = Tables.getcolumn(c, 1)
-  nn(s, v) |> first
-end
+_search(s::KDTree, x) = nn(s, Tables.matrix(x) |> vec) |> first
 
 _search((X, td), x) = pairwise(td, X, x) |> vec |> argmin

From 567bde405ffbe59d48fa2afcf3d3f5bcc499da5e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=BAlio=20Hoffimann?= <julio.hoffimann@gmail.com>
Date: Sun, 19 Jul 2026 13:21:00 -0300
Subject: [PATCH 3/3] Fix annoying random failure with elapsed time in Sample

---
 test/transforms/sample.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/transforms/sample.jl b/test/transforms/sample.jl
index 74834a83..28947760 100644
--- a/test/transforms/sample.jl
+++ b/test/transforms/sample.jl
@@ -64,5 +64,5 @@
   t = (; x, y, c)
 
   T = Sample(10_000)
-  @test @elapsed(apply(T, t)) < 0.5
+  @test @elapsed(apply(T, t)) < 1.0
 end