Fix edge case bug

SyphonArch · May 16, 2024 · 5262726 · 5262726
1 parent bc5f495
commit 5262726
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 11 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "flash1dkmeans"
-version = "0.1.0"
+version = "0.1.1"
 authors = [
   { name="Jake Hyun", email="[email protected]" },
 ]

diff --git a/src/flash1dkmeans/two_cluster.py b/src/flash1dkmeans/two_cluster.py
@@ -104,17 +104,17 @@ def numba_kmeans_1d_two_cluster(
     right_centroid = None
     division_point = None
 
-    while floor + 1 < ceiling:
+    while floor < ceiling:
         division_point = (floor + ceiling) // 2
         # If the left cluster has no weight, we need to move the floor up
         left_weight_sum = query_prefix_sum(weights_prefix_sum, start_idx, division_point)
         if left_weight_sum == 0:
-            floor = division_point
+            floor = division_point + 1
             continue
         right_weight_sum = query_prefix_sum(weights_prefix_sum, division_point, stop_idx)
         # If the right cluster has no weight, we need to move the ceiling down
         if right_weight_sum == 0:
-            ceiling = division_point
+            ceiling = division_point - 1
             continue
 
         left_centroid = query_prefix_sum(weighted_X_prefix_sum, start_idx, division_point) / left_weight_sum
@@ -126,9 +126,9 @@ def numba_kmeans_1d_two_cluster(
                 # The new division point matches the previous one, so we can stop
                 break
             else:
-                floor = division_point
+                floor = division_point + 1
         else:
-            ceiling = division_point
+            ceiling = division_point - 1
 
     # initialize variables in case the loop above does not run through
     if left_centroid is None:
@@ -187,17 +187,17 @@ def numba_kmeans_1d_two_cluster_unweighted(
     right_centroid = None
     division_point = None
 
-    while floor + 1 < ceiling:
+    while floor < ceiling:
         division_point = (floor + ceiling) // 2
         # If the left cluster has no weight, we need to move the floor up
         left_weight_sum = division_point - start_idx
         if left_weight_sum == 0:
-            floor = division_point
+            floor = division_point + 1
             continue
         right_weight_sum = stop_idx - division_point
         # If the right cluster has no weight, we need to move the ceiling down
         if right_weight_sum == 0:
-            ceiling = division_point
+            ceiling = division_point - 1
             continue
 
         left_centroid = query_prefix_sum(X_prefix_sum, start_idx, division_point) / left_weight_sum
@@ -209,9 +209,9 @@ def numba_kmeans_1d_two_cluster_unweighted(
                 # The new division point matches the previous one, so we can stop
                 break
             else:
-                floor = division_point
+                floor = division_point + 1
         else:
-            ceiling = division_point
+            ceiling = division_point - 1
 
     # initialize variables in case the loop above does not run through
     if left_centroid is None: