diff globalneighbors/distance.py @ 8:e3d6919130ca

insert via BST
author Jeff Hammel <k0scist@gmail.com>
date Sun, 25 Jun 2017 11:21:28 -0700
parents 254195d0bac2
children 638fad06e556
line wrap: on
line diff
--- a/globalneighbors/distance.py	Sun Jun 25 09:13:48 2017 -0700
+++ b/globalneighbors/distance.py	Sun Jun 25 11:21:28 2017 -0700
@@ -3,6 +3,7 @@
 """
 
 import argparse
+import bisect
 import json
 import sys
 import time
@@ -51,8 +52,19 @@
     else:
         distances.append((i, new_distance))
 
+class KeyWrapper:
+    def __init__(self, iterable, key):
+        self.it = iterable
+        self.key = key
+
+    def __getitem__(self, i):
+        return self.key(self.it[i])
+
+    def __len__(self):
+        return len(self.it)
 
 def insert_distance_bisect(distances, i, new_distance, k):
+
     if not distances:
         distances.append((i, new_distance))
         return
@@ -60,10 +72,13 @@
         if len(distances) < k:
             distances.append((i, new_distance))
         return
-    indices = [0, len(distances)]
-    while True:
-        midpoint = int((indices[-1] - indices[0])/2)
 
+    point = bisect.bisect_left(KeyWrapper(distances,
+                                          key=lambda x: x[-1]),
+                               new_distance)
+    distances.insert(point, (i, new_distance))
+    if len(distances) == k+1:
+        distances.pop()
 
 def calculate_distances(locations, r=Rearth):
     """