perf: special handling of k=1 for largest_k

seung-lab · Nov 16, 2024 · eda1498 · eda1498
1 parent 04f15b7
commit eda1498
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/cc3d/__init__.py b/cc3d/__init__.py
@@ -98,6 +98,12 @@ def largest_k(
   preserve.sort(key=lambda x: x[1])
   preserve = [ x[0] for x in preserve[-k:] ]
 
+  if k == 1:
+    cc_out = (cc_labels == preserve[0])
+    if return_N:
+      return cc_out, 1
+    return cc_out
+
   try:
     import fastremap
     cc_out = fastremap.mask_except(cc_labels, preserve, in_place=True)