perf(largest_k): use faster k-largest algorithm

seung-lab · Nov 17, 2024 · b6fe047 · b6fe047
1 parent d48f6c6
commit b6fe047
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/cc3d/__init__.py b/cc3d/__init__.py
@@ -119,9 +119,11 @@ def largest_k(
       return cc_out, 1
     return cc_out
 
-  preserve = [ (i,ct) for i,ct in enumerate(cts) if i > 0 ]
+  preserve = np.argpartition(cts[1:], len(cts) - k - 1)[-k:]
+  preserve += 1
+  preserve = [ (label,cts[label]) for label in preserve ]
   preserve.sort(key=lambda x: x[1])
-  preserve = [ x[0] for x in preserve[-k:] ]
+  preserve = [ int(x[0]) for x in preserve ]
 
   try:
     import fastremap